From 0770834cf342630bb53f7197fcfe7d6b335ed65e Mon Sep 17 00:00:00 2001 From: -LAN- Date: Tue, 27 Oct 2026 15:37:53 +0800 Subject: [PATCH 01/17] refactor(tracing): one trace pipeline behind a minimal provider interface Replace the legacy ops trace queue, the unified trace dispatcher and the per-provider SDK exporters with one pipeline and one provider contract. - dify_trace_sdk (api/providers/trace/trace-sdk) defines the provider interface: verify(config), project_url(config) and send(config, trace). A provider-neutral Trace (parent-first spans, OpenTelemetry ids) and classified errors (retryable, rejected, outcome unknown). All network access goes through an injected Egress that enforces the SSRF policy. - Each provider package implements its own send and is discovered through dify.trace_providers entry points: Langfuse, LangSmith, Opik, Weave, Arize, Phoenix, Aliyun, MLflow, Databricks and Tencent. - Producers emit typed events through a TraceScope after their records commit. tasks.trace_delivery_task assembles records within the event's workspace, feeds enterprise telemetry and the Agent collector, and sends the canonical trace to the app's current provider, retrying only that send. - Tracing secrets are write-only: the console receives configured_secrets instead of masked values, blank secrets keep their stored value, and a changed destination requires entering secrets again. - Migration 53c10126d551 keeps one trace_app_config row per app and provider and adds a unique index. - Remove core/ops, the legacy gateway and dispatch modules, and the Opik, Weave, Langfuse and Arize/Phoenix SDK dependencies; declare the runtime dependencies they used to pull in transitively. - Move the tracing tests onto the mirrored test layout and drop the suites for deleted modules; the touched tracing components use logical Tailwind classes. --- api/.env.example | 7 +- api/app_factory.py | 2 - api/configs/feature/__init__.py | 37 +- api/controllers/console/agent/tracing.py | 13 +- api/controllers/console/app/ops_trace.py | 8 +- api/core/agent/cot_agent_runner.py | 12 +- api/core/agent/fc_agent_runner.py | 5 +- .../app/apps/advanced_chat/app_generator.py | 24 +- api/core/app/apps/advanced_chat/app_runner.py | 2 +- .../advanced_chat/generate_task_pipeline.py | 49 +- api/core/app/apps/agent_app/app_generator.py | 21 +- api/core/app/apps/agent_app/app_runner.py | 5 +- .../apps/agent_app/runtime_request_builder.py | 2 +- api/core/app/apps/agent_chat/app_generator.py | 13 +- api/core/app/apps/base_app_runner.py | 2 +- api/core/app/apps/chat/app_generator.py | 13 +- api/core/app/apps/completion/app_generator.py | 18 +- api/core/app/apps/workflow/app_generator.py | 28 +- .../apps/workflow/generate_task_pipeline.py | 24 +- api/core/app/entities/app_invoke_entities.py | 25 +- .../easy_ui_based_generate_task_pipeline.py | 54 +- .../task_pipeline/message_cycle_manager.py | 2 + .../agent_tool_callback_handler.py | 17 - .../workflow_tool_callback_handler.py | 4 - api/core/llm_generator/llm_generator.py | 36 +- api/core/moderation/input_moderation.py | 28 +- api/core/ops/base_trace_instance.py | 144 - api/core/ops/entities/__init__.py | 0 api/core/ops/entities/config_entity.py | 64 - api/core/ops/entities/trace_entity.py | 292 -- api/core/ops/exceptions.py | 37 - api/core/ops/ops_trace_manager.py | 1589 --------- api/core/ops/unified_trace/__init__.py | 1 - api/core/ops/unified_trace/entities.py | 79 - api/core/ops/unified_trace/parent_context.py | 249 -- api/core/ops/unified_trace/provider.py | 82 - api/core/ops/unified_trace/registry.py | 35 - api/core/ops/unified_trace/trace_builder.py | 452 --- api/core/ops/utils.py | 175 - api/core/rag/retrieval/dataset_retrieval.py | 39 +- api/core/telemetry/__init__.py | 2 - api/core/telemetry/events.py | 51 +- api/core/telemetry/gateway.py | 44 +- api/core/tools/tool_engine.py | 23 +- .../nodes/agent_v2/runtime_request_builder.py | 2 +- api/enterprise/telemetry/enterprise_trace.py | 18 +- api/extensions/ext_application_services.py | 22 +- api/extensions/ext_celery.py | 7 + api/extensions/ext_forward_refs.py | 49 - api/extensions/ext_sentry.py | 20 - ...tore_workflow_node_execution_repository.py | 12 +- api/extensions/tracing.py | 123 + api/libs/datetime_utils.py | 20 + ...26d551_unique_trace_app_config_provider.py | 50 + api/models/model.py | 1 + api/openapi/markdown/console-openapi.md | 7 +- api/providers/trace/README.md | 115 +- .../trace/trace-aliyun/pyproject.toml | 9 +- .../src/dify_trace_aliyun/agent_trace.py | 27 - .../src/dify_trace_aliyun/aliyun_trace.py | 850 ----- .../src/dify_trace_aliyun/config.py | 102 +- .../src/dify_trace_aliyun/configuration.py | 158 - .../data_exporter/__init__.py | 0 .../data_exporter/traceclient.py | 187 -- .../dify_trace_aliyun/entities/__init__.py | 0 .../entities/aliyun_trace_entity.py | 37 - .../src/dify_trace_aliyun/entities/semconv.py | 84 - .../src/dify_trace_aliyun/provider.py | 104 + .../src/dify_trace_aliyun/utils.py | 369 --- .../data_exporter/test_traceclient.py | 160 - .../entities/test_aliyun_trace_entity.py | 88 - .../entities/test_semconv.py | 101 - .../dify_trace_aliyun/test_aliyun_trace.py | 1145 ------- .../src/dify_trace_aliyun/test_config.py | 240 +- .../src/dify_trace_aliyun/test_provider.py | 288 ++ .../src/dify_trace_aliyun/test_utils.py | 316 -- .../trace/trace-arize-phoenix/pyproject.toml | 6 +- .../_arize_project.py | 147 + .../dify_trace_arize_phoenix/agent_trace.py | 54 - .../arize_configuration.py | 304 -- .../arize_phoenix_trace.py | 1593 --------- .../src/dify_trace_arize_phoenix/config.py | 78 +- .../dify_trace_arize_phoenix/configuration.py | 182 -- .../src/dify_trace_arize_phoenix/provider.py | 285 ++ .../dify_trace_arize_phoenix/unified_trace.py | 195 -- .../test_arize_phoenix_trace.py | 2130 ------------ .../dify_trace_arize_phoenix/test_config.py | 176 +- .../dify_trace_arize_phoenix/test_provider.py | 637 ++++ .../trace/trace-langfuse/pyproject.toml | 5 +- .../src/dify_trace_langfuse/agent_trace.py | 75 - .../src/dify_trace_langfuse/config.py | 23 +- .../dify_trace_langfuse/entities/__init__.py | 0 .../entities/langfuse_trace_entity.py | 284 -- .../src/dify_trace_langfuse/langfuse_trace.py | 600 ---- .../src/dify_trace_langfuse/provider.py | 221 ++ .../src/dify_trace_langfuse/test_config.py | 43 +- .../test_langfuse_trace.py | 911 ------ .../src/dify_trace_langfuse/test_provider.py | 164 + .../trace/trace-langsmith/pyproject.toml | 4 + .../src/dify_trace_langsmith/agent_trace.py | 105 - .../src/dify_trace_langsmith/config.py | 40 +- .../dify_trace_langsmith/entities/__init__.py | 0 .../entities/langsmith_trace_entity.py | 149 - .../dify_trace_langsmith/langsmith_trace.py | 532 --- .../src/dify_trace_langsmith/provider.py | 329 ++ .../src/dify_trace_langsmith/unified_trace.py | 183 -- .../src/dify_trace_langsmith/test_config.py | 85 +- .../test_langsmith_trace.py | 743 ----- .../src/dify_trace_langsmith/test_provider.py | 376 +++ .../trace/trace-mlflow/pyproject.toml | 5 + .../src/dify_trace_mlflow/_databricks.py | 178 ++ .../src/dify_trace_mlflow/_spans.py | 203 ++ .../src/dify_trace_mlflow/agent_trace.py | 259 -- .../src/dify_trace_mlflow/config.py | 50 +- .../src/dify_trace_mlflow/configuration.py | 163 - .../src/dify_trace_mlflow/mlflow_trace.py | 694 ---- .../dify_trace_mlflow/oss_configuration.py | 201 -- .../src/dify_trace_mlflow/provider.py | 349 ++ .../src/dify_trace_mlflow/test_config.py | 60 + .../dify_trace_mlflow/test_mlflow_trace.py | 1052 ------ .../src/dify_trace_mlflow/test_provider.py | 865 +++++ api/providers/trace/trace-opik/pyproject.toml | 5 +- .../src/dify_trace_opik/agent_trace.py | 129 - .../trace-opik/src/dify_trace_opik/config.py | 40 +- .../src/dify_trace_opik/configuration.py | 156 - .../src/dify_trace_opik/opik_trace.py | 479 --- .../src/dify_trace_opik/provider.py | 334 ++ .../src/dify_trace_opik/test_config.py | 75 +- .../src/dify_trace_opik/test_opik_trace.py | 1052 ------ .../src/dify_trace_opik/test_provider.py | 526 +++ api/providers/trace/trace-sdk/pyproject.toml | 18 + .../trace-sdk/src/dify_trace_sdk/__init__.py | 5 + .../trace-sdk/src/dify_trace_sdk/config.py | 99 + .../trace-sdk/src/dify_trace_sdk/egress.py | 54 + .../trace-sdk/src/dify_trace_sdk/emitter.py | 50 + .../trace-sdk/src/dify_trace_sdk/errors.py | 47 + .../trace-sdk/src/dify_trace_sdk/events.py | 164 + .../trace-sdk/src/dify_trace_sdk/exporting.py | 81 + .../trace-sdk/src/dify_trace_sdk/genai.py | 71 + .../trace/trace-sdk/src/dify_trace_sdk/ids.py | 41 + .../trace-sdk/src/dify_trace_sdk/model.py | 194 ++ .../trace-sdk/src/dify_trace_sdk/otlp.py | 138 + .../trace-sdk/src/dify_trace_sdk/provider.py | 77 + .../trace-sdk/src/dify_trace_sdk/py.typed} | 0 .../trace-sdk/src/dify_trace_sdk/testing.py | 204 ++ .../src/dify_trace_sdk/transports.py | 67 + .../src/dify_trace_sdk/test_config.py | 51 + .../src/dify_trace_sdk/test_emitter.py | 40 + .../src/dify_trace_sdk/test_exporting.py | 49 + .../unit_tests/src/dify_trace_sdk/test_ids.py | 16 + .../src/dify_trace_sdk/test_model.py | 115 + .../src/dify_trace_sdk/test_otlp.py | 111 + .../src/dify_trace_sdk/test_testing.py | 9 + .../src/dify_trace_sdk/test_transports.py | 23 + .../trace/trace-tencent/pyproject.toml | 9 +- .../src/dify_trace_tencent/agent_trace.py | 19 - .../src/dify_trace_tencent/client.py | 569 ---- .../src/dify_trace_tencent/config.py | 77 +- .../src/dify_trace_tencent/configuration.py | 142 - .../dify_trace_tencent/entities/__init__.py | 1 - .../dify_trace_tencent/entities/semconv.py | 89 - .../entities/tencent_trace_entity.py | 21 - .../src/dify_trace_tencent/provider.py | 98 + .../src/dify_trace_tencent/span_builder.py | 383 --- .../src/dify_trace_tencent/tencent_trace.py | 538 ---- .../src/dify_trace_tencent/utils.py | 64 - .../src/dify_trace_tencent/test_client.py | 616 ---- .../src/dify_trace_tencent/test_config.py | 142 + .../src/dify_trace_tencent/test_provider.py | 302 ++ .../dify_trace_tencent/test_span_builder.py | 361 --- .../dify_trace_tencent/test_tencent_trace.py | 754 ----- .../src/dify_trace_tencent/test_utils.py | 105 - .../trace/trace-weave/pyproject.toml | 5 +- .../src/dify_trace_weave/agent_trace.py | 89 - .../src/dify_trace_weave/config.py | 46 +- .../src/dify_trace_weave/configuration.py | 250 -- .../src/dify_trace_weave/entities/__init__.py | 0 .../entities/weave_trace_entity.py | 100 - .../src/dify_trace_weave/provider.py | 374 +++ .../src/dify_trace_weave/transport.py | 56 - .../src/dify_trace_weave/weave_trace.py | 544 ---- .../src/dify_trace_weave/test_config.py | 103 +- .../src/dify_trace_weave/test_provider.py | 634 ++++ .../src/dify_trace_weave/test_weave_trace.py | 1109 ------- api/pyproject.toml | 6 + api/pytest.ini | 5 +- .../application_trace_repository.py | 51 +- api/services/app_tracing_config_gateway.py | 499 --- api/services/app_tracing_config_service.py | 14 +- api/services/conversation_name_generator.py | 2 + .../message_suggested_questions_generator.py | 26 +- api/services/tracing/assembly.py | 744 +++++ api/services/tracing/canonical.py | 444 +++ api/services/tracing/delivery.py | 271 ++ api/services/tracing/egress.py | 143 + api/services/tracing/emitter.py | 129 + api/services/tracing/genai_attributes.py | 58 - .../tracing}/hierarchy.py | 101 +- api/services/tracing/normalization.py | 251 +- api/services/tracing/otlp_export.py | 202 -- api/services/tracing/provider_common.py | 127 - api/services/tracing/provider_configs.py | 222 ++ api/services/tracing/provider_dispatch.py | 122 - api/services/tracing/records.py | 207 ++ api/services/tracing/recovery.py | 70 +- api/services/tracing/registry.py | 71 + api/services/tracing/runtime_bridge.py | 39 +- api/services/tracing/sdk_transport.py | 42 - api/services/tracing/service.py | 367 +-- api/services/workflow_persistence.py | 47 +- api/services/workflow_run_agg.py | 2 +- api/tasks/ops_trace_task.py | 156 +- api/tasks/trace_delivery_task.py | 112 + .../test_human_input_resume_node_execution.py | 1 + ...st_message_suggested_questions_failures.py | 13 +- .../app/test_message_suggested_questions.py | 5 +- .../completion/test_prepared_completion.py | 2 +- .../workflow/nodes/tool/test_tool_node.py | 19 +- .../helpers/unit/core/ops/trace_fixtures.py | 60 - .../services/conversation_name_fixtures.py | 10 - api/tests/unit_tests/.ruff.toml | 7 +- .../unit_tests/configs/test_app_config.py | 48 +- api/tests/unit_tests/conftest.py | 6 +- .../controllers/console/agent/test_tracing.py | 215 +- .../console/agent/test_tracing_runs.py | 10 +- .../controllers/console/app/test_ops_trace.py | 172 +- .../inner_api/agent/test_tracing.py | 18 +- .../service_api/app/test_workflow_events.py | 4 +- .../core/agent/test_cot_agent_runner.py | 37 +- .../core/agent/test_fc_agent_runner.py | 29 +- .../apps/advanced_chat/test_app_generator.py | 262 +- .../test_generate_task_pipeline.py | 406 +-- .../app/apps/agent_app/test_app_generator.py | 123 +- .../agent_app/test_runtime_request_builder.py | 20 +- .../app/apps/agent_chat/test_app_generator.py | 47 +- .../core/app/apps/chat/test_app_generator.py | 49 +- .../app/apps/completion/test_app_generator.py | 42 +- .../app/apps/pipeline/test_pipeline_runner.py | 2 +- .../core/app/apps/test_base_app_runner.py | 49 +- .../apps/test_message_based_app_generator.py | 2 +- .../app/apps/workflow/test_app_generator.py | 147 +- .../workflow/test_generate_task_pipeline.py | 4 +- .../app/entities/test_app_invoke_entities.py | 62 +- ...st_easy_ui_based_generate_task_pipeline.py | 259 +- .../test_message_cycle_manager.py | 7 +- .../test_agent_tool_callback_handler.py | 24 +- .../test_workflow_tool_callback_handler.py | 22 - .../core/llm_generator/test_llm_generator.py | 85 +- .../core/moderation/test_input_moderation.py | 173 +- api/tests/unit_tests/core/ops/__init__.py | 1 - .../core/ops/entities/test_config_entity.py | 75 - .../core/ops/entities/test_trace_entity.py | 168 - .../core/ops/test_base_trace_instance.py | 264 -- .../core/ops/test_ops_trace_manager.py | 1685 ---------- api/tests/unit_tests/core/ops/test_utils.py | 337 -- .../core/ops/unified_trace/__init__.py | 0 .../core/ops/unified_trace/test_entities.py | 108 - .../core/ops/unified_trace/test_hierarchy.py | 107 - .../ops/unified_trace/test_parent_context.py | 184 -- .../core/ops/unified_trace/test_provider.py | 120 - .../core/ops/unified_trace/test_registry.py | 14 - .../ops/unified_trace/test_trace_builder.py | 676 ---- .../rag/retrieval/test_dataset_retrieval.py | 149 +- .../core/telemetry/test___init__.py | 162 +- .../unit_tests/core/telemetry/test_gateway.py | 209 +- .../unit_tests/core/tools/test_tool_engine.py | 74 + .../agent_v2/test_runtime_request_builder.py | 12 +- .../core/workflow/test_node_runtime.py | 11 +- .../test_workflow_tool_container_handler.py | 1 + .../enterprise/telemetry/test_contracts.py | 4 +- .../enterprise/telemetry/test_draft_trace.py | 123 +- .../telemetry/test_enterprise_trace.py | 49 +- .../test_ext_application_services.py | 45 +- .../unit_tests/extensions/test_tracing.py | 96 + .../unit_tests/libs/test_datetime_utils.py | 23 +- ...26d551_unique_trace_app_config_provider.py | 74 + api/tests/unit_tests/models/test_model.py | 158 +- .../dify_trace_aliyun/test_configuration.py | 276 -- .../test_arize_configuration.py | 581 ---- .../test_configuration.py | 336 -- .../test_unified_trace.py | 258 -- .../test_unified_trace.py | 328 -- .../dify_trace_mlflow/test_configuration.py | 230 -- .../dify_trace_mlflow/test_mlflow_trace.py | 639 ---- .../test_oss_configuration.py | 325 -- .../src/dify_trace_opik/test_configuration.py | 188 -- .../src/dify_trace_opik/test_opik_trace.py | 214 -- .../dify_trace_tencent/test_configuration.py | 390 --- .../src/dify_trace_tencent/test_utils.py | 194 -- .../dify_trace_weave/test_configuration.py | 258 -- api/tests/unit_tests/pyrefly.toml | 7 - .../test_application_trace_repository.py | 169 +- .../test_agent_tracing_config_service.py | 5 +- .../test_app_tracing_config_gateway.py | 2810 ---------------- .../test_app_tracing_config_service.py | 18 +- .../test_conversation_name_generator.py | 8 +- ...t_message_suggested_questions_generator.py | 52 +- .../services/test_workflow_persistence.py | 154 +- .../services/test_workflow_run_index.py | 1 + .../services/tracing/test_assembly.py | 205 ++ .../services/tracing/test_canonical.py | 271 ++ .../services/tracing/test_delivery.py | 331 ++ .../services/tracing/test_egress.py | 200 ++ .../services/tracing/test_emitter.py | 202 ++ .../services/tracing/test_normalization.py | 287 +- .../services/tracing/test_otlp_export.py | 241 -- .../services/tracing/test_provider_configs.py | 272 ++ .../tracing/test_provider_dispatch.py | 2846 ----------------- .../services/tracing/test_recovery.py | 690 ++-- .../services/tracing/test_registry.py | 56 + .../tracing/test_request_log_protection.py | 33 +- .../services/tracing/test_runtime_bridge.py | 493 +-- .../services/tracing/test_service.py | 1303 ++++++-- .../tasks/test_application_trace_task.py | 55 +- .../unit_tests/tasks/test_ops_trace_task.py | 492 +-- .../tasks/test_trace_delivery_task.py | 184 ++ api/uv.lock | 524 +-- docker/envs/core-services/shared.env.example | 10 +- oxlint-suppressions.json | 10 - .../generated/api/console/agent/errors.gen.ts | 1 + .../generated/api/console/agent/orpc.gen.ts | 5 +- .../generated/api/console/agent/types.gen.ts | 3 +- .../generated/api/console/agent/zod.gen.ts | 1 + .../generated/api/console/apps/errors.gen.ts | 1 + .../generated/api/console/apps/orpc.gen.ts | 5 +- .../generated/api/console/apps/types.gen.ts | 3 +- .../generated/api/console/apps/zod.gen.ts | 1 + .../app/tracing/__tests__/panel.spec.tsx | 64 +- .../__tests__/provider-config-modal.spec.tsx | 396 ++- .../__tests__/selection.browser.spec.tsx | 14 +- .../components/app/tracing/config-popup.tsx | 4 +- .../components/app/tracing/configuration.tsx | 6 +- .../app/tracing/provider-config-modal.tsx | 674 +--- .../components/app/tracing/provider-panel.tsx | 11 +- web/app/components/app/tracing/providers.ts | 325 +- .../__tests__/tracing-configuration.spec.tsx | 42 +- web/i18n/locales/ar-TN/app.json | 1 + web/i18n/locales/az-AZ/app.json | 1 + web/i18n/locales/de-DE/app.json | 1 + web/i18n/locales/en-US/app.json | 1 + web/i18n/locales/es-ES/app.json | 1 + web/i18n/locales/fa-IR/app.json | 1 + web/i18n/locales/fr-FR/app.json | 1 + web/i18n/locales/hi-IN/app.json | 1 + web/i18n/locales/id-ID/app.json | 1 + web/i18n/locales/it-IT/app.json | 1 + web/i18n/locales/ja-JP/app.json | 1 + web/i18n/locales/ko-KR/app.json | 1 + web/i18n/locales/lo-LA/app.json | 1 + web/i18n/locales/nl-NL/app.json | 1 + web/i18n/locales/pl-PL/app.json | 1 + web/i18n/locales/pt-BR/app.json | 1 + web/i18n/locales/ro-RO/app.json | 1 + web/i18n/locales/ru-RU/app.json | 1 + web/i18n/locales/sl-SI/app.json | 1 + web/i18n/locales/th-TH/app.json | 1 + web/i18n/locales/tr-TR/app.json | 1 + web/i18n/locales/uk-UA/app.json | 1 + web/i18n/locales/vi-VN/app.json | 1 + web/i18n/locales/zh-Hans/app.json | 1 + web/i18n/locales/zh-Hant/app.json | 1 + web/i18n/resources.generated.d.ts | 1 + 362 files changed, 19660 insertions(+), 45346 deletions(-) delete mode 100644 api/core/ops/base_trace_instance.py delete mode 100644 api/core/ops/entities/__init__.py delete mode 100644 api/core/ops/entities/config_entity.py delete mode 100644 api/core/ops/entities/trace_entity.py delete mode 100644 api/core/ops/exceptions.py delete mode 100644 api/core/ops/ops_trace_manager.py delete mode 100644 api/core/ops/unified_trace/__init__.py delete mode 100644 api/core/ops/unified_trace/entities.py delete mode 100644 api/core/ops/unified_trace/parent_context.py delete mode 100644 api/core/ops/unified_trace/provider.py delete mode 100644 api/core/ops/unified_trace/registry.py delete mode 100644 api/core/ops/unified_trace/trace_builder.py delete mode 100644 api/core/ops/utils.py delete mode 100644 api/extensions/ext_forward_refs.py create mode 100644 api/extensions/tracing.py create mode 100644 api/migrations/versions/2026_10_23_1200-53c10126d551_unique_trace_app_config_provider.py delete mode 100644 api/providers/trace/trace-aliyun/src/dify_trace_aliyun/agent_trace.py delete mode 100644 api/providers/trace/trace-aliyun/src/dify_trace_aliyun/aliyun_trace.py delete mode 100644 api/providers/trace/trace-aliyun/src/dify_trace_aliyun/configuration.py delete mode 100644 api/providers/trace/trace-aliyun/src/dify_trace_aliyun/data_exporter/__init__.py delete mode 100644 api/providers/trace/trace-aliyun/src/dify_trace_aliyun/data_exporter/traceclient.py delete mode 100644 api/providers/trace/trace-aliyun/src/dify_trace_aliyun/entities/__init__.py delete mode 100644 api/providers/trace/trace-aliyun/src/dify_trace_aliyun/entities/aliyun_trace_entity.py delete mode 100644 api/providers/trace/trace-aliyun/src/dify_trace_aliyun/entities/semconv.py create mode 100644 api/providers/trace/trace-aliyun/src/dify_trace_aliyun/provider.py delete mode 100644 api/providers/trace/trace-aliyun/src/dify_trace_aliyun/utils.py delete mode 100644 api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/data_exporter/test_traceclient.py delete mode 100644 api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/entities/test_aliyun_trace_entity.py delete mode 100644 api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/entities/test_semconv.py delete mode 100644 api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_aliyun_trace.py create mode 100644 api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_provider.py delete mode 100644 api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_utils.py create mode 100644 api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/_arize_project.py delete mode 100644 api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/agent_trace.py delete mode 100644 api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/arize_configuration.py delete mode 100644 api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/arize_phoenix_trace.py delete mode 100644 api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/configuration.py create mode 100644 api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/provider.py delete mode 100644 api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/unified_trace.py delete mode 100644 api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_arize_phoenix_trace.py create mode 100644 api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_provider.py delete mode 100644 api/providers/trace/trace-langfuse/src/dify_trace_langfuse/agent_trace.py delete mode 100644 api/providers/trace/trace-langfuse/src/dify_trace_langfuse/entities/__init__.py delete mode 100644 api/providers/trace/trace-langfuse/src/dify_trace_langfuse/entities/langfuse_trace_entity.py delete mode 100644 api/providers/trace/trace-langfuse/src/dify_trace_langfuse/langfuse_trace.py create mode 100644 api/providers/trace/trace-langfuse/src/dify_trace_langfuse/provider.py delete mode 100644 api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_langfuse_trace.py create mode 100644 api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_provider.py delete mode 100644 api/providers/trace/trace-langsmith/src/dify_trace_langsmith/agent_trace.py delete mode 100644 api/providers/trace/trace-langsmith/src/dify_trace_langsmith/entities/__init__.py delete mode 100644 api/providers/trace/trace-langsmith/src/dify_trace_langsmith/entities/langsmith_trace_entity.py delete mode 100644 api/providers/trace/trace-langsmith/src/dify_trace_langsmith/langsmith_trace.py create mode 100644 api/providers/trace/trace-langsmith/src/dify_trace_langsmith/provider.py delete mode 100644 api/providers/trace/trace-langsmith/src/dify_trace_langsmith/unified_trace.py delete mode 100644 api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_langsmith_trace.py create mode 100644 api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_provider.py create mode 100644 api/providers/trace/trace-mlflow/src/dify_trace_mlflow/_databricks.py create mode 100644 api/providers/trace/trace-mlflow/src/dify_trace_mlflow/_spans.py delete mode 100644 api/providers/trace/trace-mlflow/src/dify_trace_mlflow/agent_trace.py delete mode 100644 api/providers/trace/trace-mlflow/src/dify_trace_mlflow/configuration.py delete mode 100644 api/providers/trace/trace-mlflow/src/dify_trace_mlflow/mlflow_trace.py delete mode 100644 api/providers/trace/trace-mlflow/src/dify_trace_mlflow/oss_configuration.py create mode 100644 api/providers/trace/trace-mlflow/src/dify_trace_mlflow/provider.py create mode 100644 api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_config.py delete mode 100644 api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_mlflow_trace.py create mode 100644 api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_provider.py delete mode 100644 api/providers/trace/trace-opik/src/dify_trace_opik/agent_trace.py delete mode 100644 api/providers/trace/trace-opik/src/dify_trace_opik/configuration.py delete mode 100644 api/providers/trace/trace-opik/src/dify_trace_opik/opik_trace.py create mode 100644 api/providers/trace/trace-opik/src/dify_trace_opik/provider.py delete mode 100644 api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_opik_trace.py create mode 100644 api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_provider.py create mode 100644 api/providers/trace/trace-sdk/pyproject.toml create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/__init__.py create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/config.py create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/egress.py create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/emitter.py create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/errors.py create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/events.py create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/exporting.py create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/genai.py create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/ids.py create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/model.py create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/otlp.py create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/provider.py rename api/{core/ops/__init__.py => providers/trace/trace-sdk/src/dify_trace_sdk/py.typed} (100%) create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/testing.py create mode 100644 api/providers/trace/trace-sdk/src/dify_trace_sdk/transports.py create mode 100644 api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_config.py create mode 100644 api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_emitter.py create mode 100644 api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_exporting.py create mode 100644 api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_ids.py create mode 100644 api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_model.py create mode 100644 api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_otlp.py create mode 100644 api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_testing.py create mode 100644 api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_transports.py delete mode 100644 api/providers/trace/trace-tencent/src/dify_trace_tencent/agent_trace.py delete mode 100644 api/providers/trace/trace-tencent/src/dify_trace_tencent/client.py delete mode 100644 api/providers/trace/trace-tencent/src/dify_trace_tencent/configuration.py delete mode 100644 api/providers/trace/trace-tencent/src/dify_trace_tencent/entities/__init__.py delete mode 100644 api/providers/trace/trace-tencent/src/dify_trace_tencent/entities/semconv.py delete mode 100644 api/providers/trace/trace-tencent/src/dify_trace_tencent/entities/tencent_trace_entity.py create mode 100644 api/providers/trace/trace-tencent/src/dify_trace_tencent/provider.py delete mode 100644 api/providers/trace/trace-tencent/src/dify_trace_tencent/span_builder.py delete mode 100644 api/providers/trace/trace-tencent/src/dify_trace_tencent/tencent_trace.py delete mode 100644 api/providers/trace/trace-tencent/src/dify_trace_tencent/utils.py delete mode 100644 api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_client.py create mode 100644 api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_config.py create mode 100644 api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_provider.py delete mode 100644 api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_span_builder.py delete mode 100644 api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_tencent_trace.py delete mode 100644 api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_utils.py delete mode 100644 api/providers/trace/trace-weave/src/dify_trace_weave/agent_trace.py delete mode 100644 api/providers/trace/trace-weave/src/dify_trace_weave/configuration.py delete mode 100644 api/providers/trace/trace-weave/src/dify_trace_weave/entities/__init__.py delete mode 100644 api/providers/trace/trace-weave/src/dify_trace_weave/entities/weave_trace_entity.py create mode 100644 api/providers/trace/trace-weave/src/dify_trace_weave/provider.py delete mode 100644 api/providers/trace/trace-weave/src/dify_trace_weave/transport.py delete mode 100644 api/providers/trace/trace-weave/src/dify_trace_weave/weave_trace.py create mode 100644 api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_provider.py delete mode 100644 api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_weave_trace.py delete mode 100644 api/services/app_tracing_config_gateway.py create mode 100644 api/services/tracing/assembly.py create mode 100644 api/services/tracing/canonical.py create mode 100644 api/services/tracing/delivery.py create mode 100644 api/services/tracing/egress.py create mode 100644 api/services/tracing/emitter.py delete mode 100644 api/services/tracing/genai_attributes.py rename api/{core/ops/unified_trace => services/tracing}/hierarchy.py (52%) delete mode 100644 api/services/tracing/otlp_export.py delete mode 100644 api/services/tracing/provider_common.py create mode 100644 api/services/tracing/provider_configs.py delete mode 100644 api/services/tracing/provider_dispatch.py create mode 100644 api/services/tracing/records.py create mode 100644 api/services/tracing/registry.py delete mode 100644 api/services/tracing/sdk_transport.py create mode 100644 api/tasks/trace_delivery_task.py delete mode 100644 api/tests/helpers/unit/core/ops/trace_fixtures.py delete mode 100644 api/tests/unit_tests/core/ops/__init__.py delete mode 100644 api/tests/unit_tests/core/ops/entities/test_config_entity.py delete mode 100644 api/tests/unit_tests/core/ops/entities/test_trace_entity.py delete mode 100644 api/tests/unit_tests/core/ops/test_base_trace_instance.py delete mode 100644 api/tests/unit_tests/core/ops/test_ops_trace_manager.py delete mode 100644 api/tests/unit_tests/core/ops/test_utils.py delete mode 100644 api/tests/unit_tests/core/ops/unified_trace/__init__.py delete mode 100644 api/tests/unit_tests/core/ops/unified_trace/test_entities.py delete mode 100644 api/tests/unit_tests/core/ops/unified_trace/test_hierarchy.py delete mode 100644 api/tests/unit_tests/core/ops/unified_trace/test_parent_context.py delete mode 100644 api/tests/unit_tests/core/ops/unified_trace/test_provider.py delete mode 100644 api/tests/unit_tests/core/ops/unified_trace/test_registry.py delete mode 100644 api/tests/unit_tests/core/ops/unified_trace/test_trace_builder.py create mode 100644 api/tests/unit_tests/extensions/test_tracing.py create mode 100644 api/tests/unit_tests/migrations/versions/test_2026_10_23_1200-53c10126d551_unique_trace_app_config_provider.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-aliyun/src/dify_trace_aliyun/test_configuration.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_arize_configuration.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_configuration.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_unified_trace.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-langsmith/src/dify_trace_langsmith/test_unified_trace.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_configuration.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_mlflow_trace.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_oss_configuration.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_configuration.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_opik_trace.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-tencent/src/dify_trace_tencent/test_configuration.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-tencent/src/dify_trace_tencent/test_utils.py delete mode 100644 api/tests/unit_tests/providers/trace/trace-weave/src/dify_trace_weave/test_configuration.py delete mode 100644 api/tests/unit_tests/services/test_app_tracing_config_gateway.py create mode 100644 api/tests/unit_tests/services/tracing/test_assembly.py create mode 100644 api/tests/unit_tests/services/tracing/test_canonical.py create mode 100644 api/tests/unit_tests/services/tracing/test_delivery.py create mode 100644 api/tests/unit_tests/services/tracing/test_egress.py create mode 100644 api/tests/unit_tests/services/tracing/test_emitter.py delete mode 100644 api/tests/unit_tests/services/tracing/test_otlp_export.py create mode 100644 api/tests/unit_tests/services/tracing/test_provider_configs.py delete mode 100644 api/tests/unit_tests/services/tracing/test_provider_dispatch.py create mode 100644 api/tests/unit_tests/services/tracing/test_registry.py create mode 100644 api/tests/unit_tests/tasks/test_trace_delivery_task.py diff --git a/api/.env.example b/api/.env.example index e051a8ca8f5b2d..4ed6e896513124 100644 --- a/api/.env.example +++ b/api/.env.example @@ -118,10 +118,9 @@ 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 +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/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 2ac8c78f33170c..67d4168d38f99a 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 @@ -1446,37 +1445,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( @@ -1486,6 +1459,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..3cbb97ebf2b1d7 100644 --- a/api/controllers/console/agent/tracing.py +++ b/api/controllers/console/agent/tracing.py @@ -49,6 +49,7 @@ AppTracingConfigProviderUnavailableError, AppTracingConfigVerificationFailedError, ) +from services.tracing.provider_configs import capabilities_payload _EDIT_ROLES = frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR}) @@ -206,13 +207,13 @@ 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 = capabilities_payload(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}, @@ -221,7 +222,7 @@ def _provider_capabilities(): @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 +235,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/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/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/__init__.py b/api/core/ops/entities/__init__.py deleted file mode 100644 index e69de29bb2d1d6..00000000000000 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/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/enterprise/telemetry/enterprise_trace.py b/api/enterprise/telemetry/enterprise_trace.py index 42874e1121fae7..5a9f27b5fcd8c9 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__) diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index 289f0f34ccf31c..145c46ea9057c8 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -17,10 +17,10 @@ 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 +from enterprise.telemetry.exporter import is_enterprise_telemetry_enabled from enums import DeploymentEdition, WebAppAccessMode from extensions.application_services.account import AccountServices, build_account_services from extensions.application_services.agent import AgentAppServices, build_agent_app_services @@ -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,7 +133,6 @@ 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 @@ -217,6 +217,9 @@ 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.egress import SSRFEgress +from services.tracing.emitter import CeleryTraceEmitter, TraceScopes +from services.tracing.provider_configs import ProviderConfigGateway from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks from services.upload_file_delivery_service import UploadFileDeliveryService from services.web_app_runtime_query_service import WebAppRuntimeQueryService @@ -518,7 +521,7 @@ def build_application_services( ) tracing_configs = AppTracingConfigService( configs=SQLAlchemyAppTracingConfigRepository(session_factory=database_client), - provider=OpsTraceManagerGateway(), + provider=ProviderConfigGateway(SSRFEgress()), ) apps = build_app_services( database_client=database_client, @@ -874,14 +877,23 @@ def init_app(app: Flask) -> None: app.extensions["knowledge_fs_node_dependencies"] = build_knowledge_fs_node_dependencies ToolManager.register_hardcoded_provider_class(KnowledgeFSProvider) + session_maker = get_session_maker() services = build_application_services( - database_client=get_session_maker(), + database_client=session_maker, deployment_edition=dify_config.DEPLOYMENT_EDITION, initialization_password=dify_config.INIT_PASSWORD, redis=redis_client, ) app.extensions[_EXTENSION_KEY] = services - bind_application_tracing_hooks(app, RuntimeApplicationTracingHooks()) + bind_tracing( + app, + scopes=TraceScopes( + emitter=CeleryTraceEmitter(), + session_factory=session_maker, + enterprise_enabled=is_enterprise_telemetry_enabled, + ), + hooks=RuntimeApplicationTracingHooks(), + ) 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..faccd9e5782e55 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -209,6 +209,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", # drains trace tasks enqueued before the pipeline upgrade ] day = dify_config.CELERY_BEAT_SCHEDULER_TIME @@ -254,6 +256,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_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..cb33ae906fc2cb --- /dev/null +++ b/api/extensions/tracing.py @@ -0,0 +1,123 @@ +"""Execution-side ports of the trace pipeline, bound to one Flask application. + +Core producers import only these accessors; the composition root binds the +service implementations. 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, cast + +from dify_trace_sdk.emitter import TraceScope +from flask import Flask, current_app, has_app_context + +from configs import dify_config + +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" + + +class TraceScopeFactory(Protocol): + def __call__( + self, + *, + tenant_id: str, + app_id: str | None, + app_tracing: str | None, + user_id: str | None, + extras: Mapping[str, Any], + ) -> TraceScope | None: ... + + def for_app( + self, *, tenant_id: str, app_id: str | None, user_id: str | None, extras: Mapping[str, Any] + ) -> TraceScope | None: ... + + +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: TraceScopeFactory, hooks: ApplicationTracingHooks) -> None: + app.extensions[_SCOPES_KEY] = scopes + app.extensions[_HOOKS_KEY] = hooks + + +def _scopes() -> TraceScopeFactory | None: + if not has_app_context(): + return None + return cast(TraceScopeFactory | None, current_app.extensions.get(_SCOPES_KEY)) + + +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(tenant_id=tenant_id, app_id=app_id, app_tracing=app_tracing, user_id=user_id, 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, 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, or ``None`` when Agent tracing is off or cannot run here.""" + if not dify_config.OPS_AGENT_TRACE_ENABLED: + return None + hooks = cast(ApplicationTracingHooks | None, current_app.extensions.get(_HOOKS_KEY)) if has_app_context() else None + if hooks is None: + # Agent runs started here are reported as not collected instead of failing. + logger.warning("Agent tracing is enabled but its hooks are not bound in this context") + return hooks 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..e81783b3329ff1 --- /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: b9d3f7a5c102 +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 = "b9d3f7a5c102" +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 f27dcd00daf26c..7ca5e6195a537b 100644 --- a/api/openapi/markdown/console-openapi.md +++ b/api/openapi/markdown/console-openapi.md @@ -1448,7 +1448,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 @@ -3807,7 +3807,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 @@ -35758,13 +35758,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..e896040bc739c5 100644 --- a/api/providers/trace/README.md +++ b/api/providers/trace/README.md @@ -1,77 +1,92 @@ # 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. -### 1. Config model (`BaseTracingConfig`) +Failure classes (`dify_trace_sdk.errors`): -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. +| 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. | -Fields fall into two groups used by the manager: +Error messages are logged and must never contain credentials, request bodies, or vendor response bodies. Raise them `from None`. -- **`secret_keys`** — names of fields that are **encrypted at rest** (API keys, tokens, passwords). -- **`other_keys`** — non-secret connection settings (hosts, project names, endpoints). +## Network access -List these key names in your `OpsTraceProviderConfigMap` entry so encrypt/decrypt and merge logic stay correct. +A provider never opens sockets, starts threads, or reads environment variables or files. All traffic goes through `self.egress`: -### 2. Trace instance (`BaseTraceInstance`) +- `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. -Subclass `BaseTraceInstance` and implement: +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. -```python -def trace(self, trace_info: BaseTraceInfo) -> None: ... -``` +## Configuration and secrets -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: +Stored configurations keep the provider's fields at the top level, with `secret()` fields encrypted per workspace. Secrets are write-only: -- `WorkflowTraceInfo`, `WorkflowNodeTraceInfo`, `DraftNodeExecutionTrace` -- `MessageTraceInfo`, `ToolTraceInfo`, `ModerationTraceInfo`, `SuggestedQuestionTraceInfo` -- `DatasetRetrievalTraceInfo`, `GenerateNameTraceInfo`, `PromptGenerationTraceInfo` +- 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. -You may ignore categories your backend does not support; existing providers often no-op unhandled types. +`ProviderConfig` hides inputs from validation errors; keep validator messages fixed texts too. -Optional: use `get_service_account_with_tenant(app_id)` from the base class when you need tenant-scoped account context. +## Adding a provider -### 3. Register in the API core +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`. -Upstream changes are required so Dify knows your provider exists: + List `py.typed` under `[tool.setuptools.package-data]`. +2. Register the provider class as an entry point: -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`. + ```toml + [project.entry-points."dify.trace_providers"] + = "dify_trace_.provider:Provider" + ``` -If the `match` case is missing, the provider string will not resolve and tracing will be disabled for that app. +3. Add the ID to `SUPPORTED_PROVIDERS` in `api/services/tracing/registry.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. -## Package layout +`trace-langfuse/` is the smallest reference implementation; `trace-arize-phoenix/` and `trace-tencent/` show OTLP over HTTP and gRPC. -Each provider is a normal uv workspace member, for example: +## How traces reach a provider -- `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`. +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. -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/`. + A scope exists only when the app has tracing enabled or enterprise telemetry is on. +2. `tasks.trace_delivery_task` delivers each event on the `ops_trace` queue: + - `services.tracing.assembly` reloads the records through the event's workspace. + - Enterprise telemetry receives every record. + - The Agent trace collector 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`. 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/__init__.py b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/data_exporter/__init__.py deleted file mode 100644 index e69de29bb2d1d6..00000000000000 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..ff05d84411f9be 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,16 @@ -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() @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 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..b587bba7c8d837 --- /dev/null +++ b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/provider.py @@ -0,0 +1,221 @@ +"""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 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() + + +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 | None: + host = config.host.rstrip("/") + return f"{host}/project/{config.project_key}" if config.project_key else 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)}/traces/{quote(trace_key, safe='')}" if config.project_key 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": root.inputs, + "output": root.outputs, + "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": span.inputs, + "output": span.outputs, + "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..d836f678671f14 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,9 @@ 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"} + 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") 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..0c04bf5edb88cc --- /dev/null +++ b/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_provider.py @@ -0,0 +1,164 @@ +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 + +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() + + +@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 result.trace_id == MESSAGE_ID + assert result.trace_url == f"https://cloud.langfuse.com/project/project-1/traces/{MESSAGE_ID}" + + +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/core/ops/__init__.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/py.typed similarity index 100% rename from api/core/ops/__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 bffd16afc1b234..4affbe879e9f65 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/application_trace_repository.py b/api/repositories/application_trace_repository.py index 653e66ba6e5c48..689b2a2e91a976 100644 --- a/api/repositories/application_trace_repository.py +++ b/api/repositories/application_trace_repository.py @@ -7,7 +7,9 @@ from uuid import UUID, uuid5 from dify_agent.protocol.tracing import AgentTraceBatch +from dify_trace_sdk.ids import span_id_for from sqlalchemy import delete, or_, select +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session, sessionmaker from libs.datetime_utils import naive_utc_now @@ -16,7 +18,7 @@ from models.model import App, Message, TraceAppConfig from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowRun from repositories.factory import DifyAPIRepositoryFactory -from services.tracing.normalization import safe_attributes, span_id +from services.tracing.normalization import safe_attributes class TraceCollectionError(Exception): @@ -81,10 +83,8 @@ def start( 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()) + app = session.scalar(select(App).where(App.id == app_id, App.tenant_id == tenant_id)) if app is None: return None existing = session.scalar( @@ -144,15 +144,13 @@ def start( config_revision=revision, provider_config=frozen, trace_id=uuidv7().hex, - root_span_id=span_id(source_id), + root_span_id=span_id_for(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", ) - session.add(record) - session.flush() - return record + return self._insert_or_get(session, record) def source_run( self, *, tenant_id: str, app_id: str, source_kind: str, source_id: str @@ -180,7 +178,7 @@ def record_initialization_failure( 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()) + app = session.scalar(select(App).where(App.id == app_id, App.tenant_id == tenant_id)) if app is None: return None current = session.scalar( @@ -203,7 +201,7 @@ def record_initialization_failure( config_revision="", provider_config={}, trace_id=uuidv7().hex, - root_span_id=span_id(source_id), + root_span_id=span_id_for(source_id), session_id=source_id, capture_content=False, capture_agent=False, @@ -211,8 +209,33 @@ def record_initialization_failure( error_code="tracing_initialization_failed", terminal_at=naive_utc_now(), ) - session.add(current) - return current + return self._insert_or_get(session, current) + + @staticmethod + def _insert_or_get(session: Session, record: ApplicationTraceRun) -> ApplicationTraceRun: + """Insert a run; when a concurrent start won the unique source key, return its run instead. + + The first frozen destination wins, so one execution never exports to two providers. + """ + try: + with session.begin_nested(): + session.add(record) + except IntegrityError: + # A locking read sees the committed winner even under MySQL's REPEATABLE READ snapshot. + existing = session.scalar( + select(ApplicationTraceRun) + .where( + ApplicationTraceRun.tenant_id == record.tenant_id, + ApplicationTraceRun.app_id == record.app_id, + ApplicationTraceRun.source_kind == record.source_kind, + ApplicationTraceRun.source_id == record.source_id, + ) + .with_for_update() + ) + if existing is None: + raise + return existing + return record def invocation( self, *, run_id: str, agent_id: str, node_execution_id: str | None, attempt: int, call_key: str @@ -234,8 +257,8 @@ def invocation( 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}"), + parent_span_id=span_id_for(node_execution_id) if node_execution_id else run.root_span_id, + span_id=span_id_for(f"agent-attempt:{identifier}"), node_execution_id=node_execution_id, attempt=attempt, ) 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..7a8a86f9475f21 100644 --- a/api/services/conversation_name_generator.py +++ b/api/services/conversation_name_generator.py @@ -8,6 +8,7 @@ from core.model_context import use_credit_usage_metadata from core.model_manager import ModelManager from extensions.ext_database import db +from extensions.tracing import create_trace_scope_for_app from graphon.model_runtime.entities.model_entities import ModelType @@ -30,4 +31,5 @@ def generate_conversation_name(*, tenant_id: str, app_id: str, conversation_id: conversation_id=conversation_id, query=query, model_instance=model, + trace_scope=create_trace_scope_for_app(tenant_id=tenant_id, app_id=app_id, user_id=None, extras={}), ) diff --git a/api/services/message_suggested_questions_generator.py b/api/services/message_suggested_questions_generator.py index 696e0aa8d435a1..bc791e72f9ee28 100644 --- a/api/services/message_suggested_questions_generator.py +++ b/api/services/message_suggested_questions_generator.py @@ -7,21 +7,21 @@ 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 extensions.tracing import create_trace_scope_for_app 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 if TYPE_CHECKING: @@ -122,7 +122,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 +131,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 = create_trace_scope_for_app( + tenant_id=context.tenant_id, app_id=context.app_id, user_id=None, extras={} ) + 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/assembly.py b/api/services/tracing/assembly.py new file mode 100644 index 00000000000000..1d42426f3c6b93 --- /dev/null +++ b/api/services/tracing/assembly.py @@ -0,0 +1,744 @@ +"""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. +""" + +import json +from collections.abc import Mapping, Sequence +from datetime import datetime, timedelta +from typing import Any + +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 core.repositories import DifyCoreRepositoryFactory +from graphon.entities import WorkflowNodeExecution +from graphon.enums import WorkflowExecutionStatus +from models import Account, 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, WorkflowNodeExecutionTriggeredFrom, WorkflowRun +from repositories.factory import DifyAPIRepositoryFactory +from services.tracing.records import ( + BaseTraceInfo, + DatasetRetrievalTraceInfo, + DraftNodeExecutionTrace, + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + PromptGenerationTraceInfo, + SuggestedQuestionTraceInfo, + ToolTraceInfo, + 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, +} + + +class SourceNotFoundError(Exception): + """The referenced record does not exist for the event's tenant and app.""" + + +class SourceNotReadyError(Exception): + """The referenced record is not persisted or not finished yet.""" + + +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 _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 TraceAssembler: + def __init__(self, session_factory: sessionmaker[Session], *, files_url: str, enterprise_details: bool) -> None: + self._sessions = session_factory + self._files_url = files_url + # App, workspace, and credential names are enterprise telemetry dimensions only. + self._enterprise_details = enterprise_details + + def assemble(self, event: TraceEvent) -> BaseTraceInfo: + """Raises SourceNotFoundError or SourceNotReadyError.""" + 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 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=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 + runs = DifyAPIRepositoryFactory.create_api_workflow_run_repository(self._sessions) + run = 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 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 = 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, + workflow_data=run, + node_executions=executions, + 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) -> Sequence[WorkflowNodeExecution]: + # Repositories take an identity only to attribute writes. This read needs no member of the + # workspace, so a run of an app whose creator left it is still traced. + reader = Account(name="Trace assembly", email="trace-assembly@localhost") + reader.id = run.created_by + repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( + session_factory=self._sessions, + tenant_id=run.tenant_id, + user=reader, + app_id=run.app_id, + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + return repository.get_by_workflow_execution(workflow_execution_id=run.id) + + @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 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=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 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=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 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=message, + error=None, + ) + + def _tool(self, event: TraceEvent, payload: ToolInvoked) -> ToolTraceInfo: + with 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=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 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=message, + 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 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 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/services/tracing/canonical.py b/api/services/tracing/canonical.py new file mode 100644 index 00000000000000..0283af7a91885d --- /dev/null +++ b/api/services/tracing/canonical.py @@ -0,0 +1,444 @@ +"""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 core.app.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: + run = record.workflow_data + executions = record.node_executions + hierarchy = build_workflow_hierarchy(executions) + workflow_start = run.created_at + workflow_end = max(run.finished_at or record.end_time or workflow_start, 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 or workflow_end, 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/delivery.py b/api/services/tracing/delivery.py new file mode 100644 index 00000000000000..592e9158ef6301 --- /dev/null +++ b/api/services/tracing/delivery.py @@ -0,0 +1,271 @@ +"""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 json +import logging +from collections.abc import Callable +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Any, Protocol +from uuid import uuid4 + +from celery import current_app as celery_app +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 sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from configs import dify_config +from enterprise.telemetry.enterprise_trace import EnterpriseOtelTrace +from enterprise.telemetry.exporter import is_enterprise_telemetry_enabled +from extensions.ext_database import db +from extensions.ext_storage import storage +from models.model import App, TraceAppConfig +from services.tracing.assembly import TraceAssembler +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.egress import SSRFEgress +from services.tracing.emitter import TRACE_QUEUE, payload_key +from services.tracing.provider_configs import ProviderConfigUnreadableError, decode_config +from services.tracing.records import BaseTraceInfo, MessageTraceInfo, WorkflowTraceInfo +from services.tracing.registry import TraceProviderUnavailableError, UnknownTraceProviderError, provider_class +from services.tracing.service import application_tracing_service + +logger = logging.getLogger(__name__) + +RETRY_TASK = "tasks.trace_delivery_task.retry_trace_export" + +_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 RecordAssembler(Protocol): + def assemble(self, event: TraceEvent) -> BaseTraceInfo: + """Raises SourceNotFoundError, or SourceNotReadyError while the source is not persisted yet.""" + ... + + +class TraceCollector(Protocol): + def capture_terminal(self, info: BaseTraceInfo) -> bool: + """Take ownership of a record of an Agent-managed execution; ``False`` leaves it to the provider.""" + ... + + def link_managed_parent(self, info: WorkflowTraceInfo) -> WorkflowTraceInfo: + """Name the Agent-collected parent that a provider-exported nested workflow cannot attach to.""" + ... + + +class TraceCollectorUnavailableError(Exception): + """The Agent trace collector could not decide ownership; the event must be redelivered.""" + + +@dataclass(frozen=True) +class Destination: + provider: TraceProvider[Any] + config: ProviderConfig + + +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) + + +def retry_delay_seconds(attempt: int) -> int: + return min(600, 10 * 2 ** (attempt - 1)) + + +class TraceDeliveryService: + def __init__( + self, + *, + session_factory: sessionmaker[Session], + assembler: RecordAssembler, + builder: CanonicalTraceBuilder, + egress: Egress, + enterprise_sink: Callable[[BaseTraceInfo], None] | None, + collector: TraceCollector | None, + max_export_attempts: int, + ) -> None: + self._sessions = session_factory + self._assembler = assembler + self._builder = builder + self._egress = egress + 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._assembler.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(record): + _event_counter.add(1, {"kind": kind, "outcome": "collected"}) + return + if isinstance(record, WorkflowTraceInfo): + record = self._collector.link_managed_parent(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._schedule_retry(job.model_copy(update={"attempt": job.attempt + 1})) + 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.""" + with self._sessions() as session: + app_tracing = session.scalar(select(App.tracing).where(App.id == app_id, App.tenant_id == tenant_id)) + provider_id = self._enabled_provider(app_tracing) + if provider_id is None: + return None + stored = session.scalar( + select(TraceAppConfig.tracing_config).where( + TraceAppConfig.app_id == app_id, + TraceAppConfig.tracing_provider == provider_id, + TraceAppConfig.is_active.is_(True), + ) + ) + if not isinstance(stored, dict): + return None + try: + provider = provider_class(provider_id) + config = decode_config(tenant_id, provider, stored) + except (UnknownTraceProviderError, TraceProviderUnavailableError): + logger.warning("Tracing provider unavailable: provider=%s app_id=%s", provider_id, app_id) + return None + except ProviderConfigUnreadableError: + logger.warning("Tracing configuration unreadable: provider=%s app_id=%s", provider_id, app_id) + return None + return Destination(provider=provider(self._egress), config=config) + + @staticmethod + def _enabled_provider(app_tracing: str | None) -> str | None: + if not app_tracing: + return None + try: + policy = json.loads(app_tracing) + except ValueError: + return None + if not isinstance(policy, dict) or policy.get("enabled") is not True: + return None + provider_id = policy.get("tracing_provider") + return provider_id if isinstance(provider_id, str) else None + + @staticmethod + def _schedule_retry(job: ExportJob) -> 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,), queue=TRACE_QUEUE, countdown=retry_delay_seconds(job.attempt - 1)) + + +def _enterprise_sink(record: BaseTraceInfo) -> None: + EnterpriseOtelTrace().trace(record) + + +def trace_delivery_service() -> TraceDeliveryService: + sessions = sessionmaker(bind=db.engine, expire_on_commit=False) + enterprise = is_enterprise_telemetry_enabled() + return TraceDeliveryService( + session_factory=sessions, + assembler=TraceAssembler(sessions, files_url=dify_config.FILES_URL, enterprise_details=enterprise), + builder=CanonicalTraceBuilder(), + egress=SSRFEgress(), + enterprise_sink=_enterprise_sink if enterprise else None, + collector=application_tracing_service() if dify_config.OPS_AGENT_TRACE_ENABLED else None, + max_export_attempts=dify_config.OPS_TRACE_EXPORT_MAX_ATTEMPTS, + ) 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/emitter.py b/api/services/tracing/emitter.py new file mode 100644 index 00000000000000..75fb7074c5b116 --- /dev/null +++ b/api/services/tracing/emitter.py @@ -0,0 +1,129 @@ +"""The request-path side of the trace pipeline. + +Deciding whether to trace costs nothing extra: generators already hold the +app's tracing policy, so no configuration is decrypted and no provider is +loaded while a user waits. Events go straight to the delivery worker; large +ones travel through storage instead of the broker. +""" + +import json +import logging +from collections.abc import Callable, Mapping +from datetime import datetime +from typing import Any, override + +from celery import current_app as celery_app +from dify_trace_sdk.emitter import TraceEmitter, TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from core.helper.trace_id_helper import ParentTraceContext +from extensions.ext_storage import storage +from models.model import App + +logger = logging.getLogger(__name__) + +DELIVER_TASK = "tasks.trace_delivery_task.deliver_trace_event" +TRACE_QUEUE = "ops_trace" +PAYLOAD_PREFIX = "trace_payloads" +INLINE_EVENT_BYTES = 256 * 1024 +# Node executions may be saved by another worker after their run row; give them time to land. +WORKFLOW_SETTLE_SECONDS = 5 + + +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" + + +def tracing_policy_enabled(app_tracing: str | None) -> bool: + if not app_tracing: + return False + try: + policy = json.loads(app_tracing) + except ValueError: + return False + return ( + isinstance(policy, dict) and policy.get("enabled") is True and isinstance(policy.get("tracing_provider"), str) + ) + + +def correlation_from_extras(extras: Mapping[str, Any]) -> TraceCorrelation: + 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, + ) + + +class CeleryTraceEmitter(TraceEmitter): + @override + def emit(self, event: TraceEvent) -> None: + encoded = event.model_dump_json() + if len(encoded.encode()) > INLINE_EVENT_BYTES: + key = payload_key(event.occurred_at, "events", event.event_id) + storage.save(key, encoded.encode()) + envelope = {"payload_key": key} + else: + envelope = {"event": encoded} + countdown = WORKFLOW_SETTLE_SECONDS if event.payload.kind == "workflow_run" else 0 + celery_app.send_task(DELIVER_TASK, args=(envelope,), queue=TRACE_QUEUE, countdown=countdown) + + +class TraceScopes: + """Create trace scopes for executions whose events someone consumes.""" + + def __init__( + self, + *, + emitter: TraceEmitter, + session_factory: sessionmaker[Session], + enterprise_enabled: Callable[[], bool], + ) -> None: + self._emitter = emitter + self._sessions = session_factory + self._enterprise_enabled = enterprise_enabled + + def __call__( + self, + *, + tenant_id: str, + app_id: str | None, + app_tracing: str | None, + user_id: str | None, + extras: Mapping[str, Any], + ) -> TraceScope | None: + if not (tracing_policy_enabled(app_tracing) 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_from_extras(extras), + ) + + def for_app( + self, *, tenant_id: str, app_id: str | None, user_id: str | None, extras: Mapping[str, Any] + ) -> TraceScope | None: + """A scope for callers that do not hold the app row; tracing failures stay out of the request.""" + app_tracing = None + if app_id is not None: + try: + with self._sessions() as session: + app_tracing = session.scalar( + select(App.tracing).where(App.id == app_id, App.tenant_id == tenant_id) + ) + except Exception: + logger.warning("Tracing policy lookup failed: app_id=%s", app_id, exc_info=True) + return None + return self(tenant_id=tenant_id, app_id=app_id, app_tracing=app_tracing, user_id=user_id, extras=extras) 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 52% rename from api/core/ops/unified_trace/hierarchy.py rename to api/services/tracing/hierarchy.py index 0f08ceda597d24..a42885b965148a 100644 --- a/api/core/ops/unified_trace/hierarchy.py +++ b/api/services/tracing/hierarchy.py @@ -1,4 +1,10 @@ -"""Deterministically reconstruct provider-neutral workflow execution hierarchy.""" +"""Deterministically reconstruct the span hierarchy of one workflow run. + +Parents come from graph predecessors and from iteration/loop containers; +each iteration or loop round gets a synthetic wrapper. Ambiguous repeated node +IDs and cyclic predecessor data are dropped rather than guessed, so those spans +fall back to the workflow root. +""" import re from collections.abc import Mapping, Sequence @@ -7,6 +13,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 +22,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 +47,29 @@ 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 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 finished_at(execution: WorkflowNodeExecution) -> datetime: + return execution.finished_at or execution.created_at + timedelta(seconds=execution.elapsed_time or 0) -def _unique_execution_by_node_id(executions: Sequence[WorkflowExecutionLike]) -> dict[str, str]: +def _unique_execution_by_node_id(executions: Sequence[WorkflowNodeExecution]) -> dict[str, str]: result: dict[str, str] = {} ambiguous: set[str] = set() for item in executions: - node_id = item.node_id - if not isinstance(node_id, str) or node_id in ambiguous: + if item.node_id in ambiguous: continue - item_execution_id = execution_id(item) - previous = result.get(node_id) + previous = result.get(item.node_id) if previous is None: - result[node_id] = item_execution_id - elif previous != item_execution_id: - result.pop(node_id, None) - ambiguous.add(node_id) + result[item.node_id] = item.id + elif previous != item.id: + result.pop(item.node_id, None) + ambiguous.add(item.node_id) return result -def _metadata_or_attr(execution: object, metadata: Mapping[str, Any], key: str) -> Any: - value = metadata.get(key) - return value if value is not None else _read_attribute(execution, key) - - def _normalize_index(value: Any) -> str | None: if isinstance(value, bool): return None @@ -95,19 +80,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 +97,37 @@ 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. - """ +def build_workflow_hierarchy(executions: Sequence[WorkflowNodeExecution]) -> WorkflowHierarchy: execution_by_node_id = _unique_execution_by_node_id(executions) parent_by_execution_id: dict[str, str] = {} for item in executions: - item_execution_id = execution_id(item) - predecessor_node_id = item.predecessor_node_id parent_execution_id = ( - execution_by_node_id.get(predecessor_node_id) if isinstance(predecessor_node_id, str) else None + execution_by_node_id.get(item.predecessor_node_id) 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) 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 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): + 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 +135,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/normalization.py b/api/services/tracing/normalization.py index 0216e12431e29b..aba79110d456ff 100644 --- a/api/services/tracing/normalization.py +++ b/api/services/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/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..d76285ee1f28d5 --- /dev/null +++ b/api/services/tracing/provider_configs.py @@ -0,0 +1,222 @@ +"""Store, verify, and present tracing provider configurations. + +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 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 core.helper.encrypter import batch_decrypt_token, encrypt_token, obfuscated_token +from services.app_tracing_config_service import ( + AppTracingConfigInvalidConfigurationError, + AppTracingConfigInvalidProviderError, + AppTracingConfigProcessingError, + AppTracingConfigProviderUnavailableError, + AppTracingConfigVerificationFailedError, + PresentedTracingConfig, + TracingConfigProviderGateway, +) +from services.tracing.registry import ( + SUPPORTED_PROVIDERS, + TraceProviderUnavailableError, + UnknownTraceProviderError, + provider_class, +) + +logger = logging.getLogger(__name__) + + +class ProviderConfigUnreadableError(Exception): + """A stored configuration cannot be decrypted or no longer validates.""" + + +def resolve_provider(provider_id: str) -> type[TraceProvider[Any]]: + try: + return provider_class(provider_id) + except UnknownTraceProviderError: + raise AppTracingConfigInvalidProviderError(provider_id) from None + except TraceProviderUnavailableError: + raise AppTracingConfigProviderUnavailableError from None + + +def encode_config(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] = encrypt_token(tenant_id, value) + return stored + + +def decode_config(tenant_id: str, provider: type[TraceProvider[Any]], stored: Mapping[str, Any]) -> ProviderConfig: + """Decrypt and validate a stored configuration for verification or delivery.""" + 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 = batch_decrypt_token(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 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)) + + +def _comparable(value: object) -> object: + return value.strip().rstrip("/") if isinstance(value, str) else value + + +def _left_unchanged(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 == obfuscated_token(current_value) + ) + + +class ProviderConfigGateway(TracingConfigProviderGateway): + def __init__(self, egress: Egress) -> None: + 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(tracing_provider) + try: + decode_config(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(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(tracing_provider) + current = None + if current_tracing_config is not None: + try: + current = decode_config(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(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)) + + @staticmethod + def _submitted( + 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 _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 _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 encode_config(workspace_id, verified) + + +def capabilities_payload(provider_id: str) -> dict[str, bool | str]: + """What a provider's mapping supports, for the console. Describes the adapter, not delivery health.""" + capabilities = resolve_provider(provider_id).capabilities + return { + "provider": provider_id, + "agent_spans": True, + "session": True, + "native_usage": True, + "native_cost": capabilities.native_cost, + "trace_url": capabilities.trace_url, + "span_url": capabilities.span_url, + "export_ack": "accepted", + } 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/records.py b/api/services/tracing/records.py new file mode 100644 index 00000000000000..72271371995598 --- /dev/null +++ b/api/services/tracing/records.py @@ -0,0 +1,207 @@ +"""Assembled trace records: one Dify operation with the data its sinks need. + +Records live only inside the delivery worker. They may reference ORM rows +because they never cross a process or storage boundary; the canonical trace +built from them is what gets persisted for retries. +""" + +from collections.abc import Mapping, Sequence +from datetime import datetime +from enum import StrEnum +from typing import Any + +from pydantic import BaseModel, ConfigDict + +from core.helper.trace_id_helper import ParentTraceContext +from graphon.entities import WorkflowNodeExecution +from models.model import Message +from models.workflow import WorkflowRun + + +class BaseTraceInfo(BaseModel): + message_id: str | None = None + message_data: Message | 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 + """The caller-supplied trace ID, if any.""" + + model_config = ConfigDict(arbitrary_types_allowed=True, 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.""" + match self.metadata.get("parent_trace_context"): + case ParentTraceContext() as context: + pass + case Mapping() as raw: + try: + context = ParentTraceContext.model_validate(raw) + except ValueError: + return None, None + case _: + return None, None + return context.parent_workflow_run_id, context.parent_node_execution_id + + +class WorkflowTraceInfo(BaseTraceInfo): + workflow_data: WorkflowRun + node_executions: Sequence[WorkflowNodeExecution] + 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: Message + 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..26943a3c4659bb 100644 --- a/api/services/tracing/recovery.py +++ b/api/services/tracing/recovery.py @@ -8,22 +8,24 @@ import json import logging import math -from datetime import datetime, timedelta -from typing import Protocol, cast +from datetime import UTC, datetime, timedelta +from typing import Protocol +from uuid import uuid4 +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent, TracePayload, WorkflowRunCompleted from sqlalchemy import select from sqlalchemy.exc import SQLAlchemyError 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 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.model import App, AppMode, Conversation, Message from models.workflow import WorkflowRun from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError +from services.tracing.assembly import SourceNotFoundError, SourceNotReadyError, TraceAssembler +from services.tracing.records import BaseTraceInfo, MessageTraceInfo, WorkflowTraceInfo logger = logging.getLogger(__name__) _SETTLE_SECONDS = 5 @@ -38,6 +40,9 @@ class RecoveryTarget(Protocol): repository: ApplicationTraceRepository + @property + def assembler(self) -> TraceAssembler: ... + def capture_terminal(self, info: BaseTraceInfo) -> bool: ... @@ -72,7 +77,24 @@ def _complete_message_metadata(message: Message) -> bool: return isinstance(latency, (int, float)) and math.isfinite(latency) and latency >= 0 -def _message_task(repository: ApplicationTraceRepository, run: ApplicationTraceRun, now: datetime) -> TraceTask | None: +def _event(run: ApplicationTraceRun, 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_task(repository: ApplicationTraceRepository, run: ApplicationTraceRun, now: datetime) -> TraceEvent | None: with repository.sessions() as session: source = session.execute( select(Message, App.mode) @@ -123,15 +145,12 @@ def _message_task(repository: ApplicationTraceRepository, run: ApplicationTraceR 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, - ) + return _event(run, MessageCompleted(message_id=message.id, conversation_id=message.conversation_id), None) -def _workflow_task(repository: ApplicationTraceRepository, run: ApplicationTraceRun, now: datetime) -> TraceTask | None: +def _workflow_task( + repository: ApplicationTraceRepository, run: ApplicationTraceRun, now: datetime +) -> TraceEvent | None: workflow = repository._source(tenant_id=run.tenant_id, app_id=run.app_id, kind="workflow", source_id=run.source_id) # An asynchronous workflow repository may not have saved the reserved run yet. if ( @@ -151,12 +170,10 @@ def _workflow_task(repository: ApplicationTraceRepository, run: ApplicationTrace ) .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, + return _event( + run, + WorkflowRunCompleted(workflow_run_id=workflow.id, conversation_id=conversation_id, total_tokens=None), + workflow.created_by, ) @@ -170,16 +187,17 @@ def recover_terminal(target: RecoveryTarget, run: ApplicationTraceRun) -> bool: return False error_code: str | None = None try: - task = (_message_task if run.source_kind == "message" else _workflow_task)( + event = (_message_task if run.source_kind == "message" else _workflow_task)( target.repository, run, naive_utc_now() ) - if task is None: + if event is None: + return False + try: + info = target.assembler.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): diff --git a/api/services/tracing/registry.py b/api/services/tracing/registry.py new file mode 100644 index 00000000000000..988b11211fdabe --- /dev/null +++ b/api/services/tracing/registry.py @@ -0,0 +1,71 @@ +"""Discover installed trace providers through ``dify.trace_providers`` entry points. + +Dify supports a fixed set of provider IDs; a distribution may install any +subset of their packages. An unknown ID is a client error, a known but +uninstalled one is a deployment fact. Neither ever raises into a business +request: callers translate both. +""" + +import functools +from importlib.metadata import EntryPoint, entry_points +from typing import Any + +from dify_trace_sdk.provider import TraceProvider + +ENTRY_POINT_GROUP = "dify.trace_providers" + +SUPPORTED_PROVIDERS = ( + "langfuse", + "langsmith", + "opik", + "weave", + "arize", + "phoenix", + "aliyun", + "mlflow", + "databricks", + "tencent", +) + + +class UnknownTraceProviderError(Exception): + def __init__(self, provider_id: str) -> None: + super().__init__(f"Unknown tracing provider: {provider_id}") + + +class TraceProviderUnavailableError(Exception): + def __init__(self, provider_id: str) -> None: + super().__init__(f"Tracing provider {provider_id} is not installed in this deployment") + + +@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 provider_class(provider_id: str) -> type[TraceProvider[Any]]: + if provider_id not in SUPPORTED_PROVIDERS: + raise UnknownTraceProviderError(provider_id) + 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 + + +def installed_providers() -> tuple[str, ...]: + """Supported providers whose packages load in this deployment, in display order.""" + installed: list[str] = [] + for provider_id in SUPPORTED_PROVIDERS: + try: + provider_class(provider_id) + except TraceProviderUnavailableError: + continue + installed.append(provider_id) + return tuple(installed) diff --git a/api/services/tracing/runtime_bridge.py b/api/services/tracing/runtime_bridge.py index f94aae3b192b8d..fec0ac560caf83 100644 --- a/api/services/tracing/runtime_bridge.py +++ b/api/services/tracing/runtime_bridge.py @@ -1,4 +1,4 @@ -"""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. @@ -18,7 +18,6 @@ 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 @@ -158,40 +157,6 @@ def freeze_workflow_run( _log_failure("workflow_start", error) -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.""" @@ -200,6 +165,4 @@ class RuntimeApplicationTracingHooks: 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) 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..c1aec1f470d5cc 100644 --- a/api/services/tracing/service.py +++ b/api/services/tracing/service.py @@ -1,4 +1,10 @@ -"""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 @@ -6,31 +12,44 @@ from collections.abc import Generator from contextlib import contextmanager from datetime import UTC, datetime, timedelta -from typing import Any, Literal, cast +from typing import Any, Literal +from uuid import uuid4 from dify_agent.layers.execution_context import ApplicationTracingContext from dify_agent.protocol.tracing import AgentTraceBatch +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.ids import span_id_for +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace +from dify_trace_sdk.provider import TraceProvider +from pydantic import ValidationError 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 import App, AppMode, 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 services.app_tracing_config_service import ( + AppTracingConfigInvalidProviderError, + AppTracingConfigProviderUnavailableError, +) +from services.tracing.assembly import TraceAssembler +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.egress import SSRFEgress +from services.tracing.normalization import ( + legacy_outer_trace, + normalize_agent_span, + order_spans, + remap_outer_trace, + sanitize, ) +from services.tracing.provider_configs import ProviderConfigUnreadableError, capabilities_payload, decode_config +from services.tracing.records import BaseTraceInfo, MessageTraceInfo, WorkflowTraceInfo +from services.tracing.recovery import agent_stop_requested, recover_terminal +from services.tracing.registry import TraceProviderUnavailableError, UnknownTraceProviderError, provider_class from services.tracing.runtime_bridge import workflow_has_agent_v2 logger = logging.getLogger(__name__) @@ -48,8 +67,18 @@ def _message_inputs_are_empty(inputs: object) -> bool: class ApplicationTracingService: - def __init__(self, repository: ApplicationTraceRepository) -> None: + def __init__( + self, + repository: ApplicationTraceRepository, + *, + assembler: TraceAssembler, + builder: CanonicalTraceBuilder, + egress: Egress, + ) -> None: self.repository = repository + self.assembler = assembler + self._builder = builder + self._egress = egress def start( self, @@ -137,39 +166,48 @@ def capture_agent_app_terminal( 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( 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": 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.assembler.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): + inputs = {"query": info.message_data.query} + if output is not None: # 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) + self.capture_terminal( + info.model_copy( + update={ + "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}", + } + ) + ) @staticmethod def schedule(run_id: str) -> None: @@ -180,19 +218,6 @@ def schedule(run_id: str) -> None: 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 _run_for_source( self, *, tenant_id: str, app_id: str, source_kind: str, source_id: str ) -> ApplicationTraceRun | None: @@ -217,7 +242,7 @@ def _run_for_source( 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.""" + """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 @@ -255,73 +280,16 @@ def _managed_owner(self, app: App, message: Message | None = None, workflow_run_ 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 + def _app(self, app_id: str) -> App | 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) - return False + return session.scalar(select(App).where(App.id == app_id)) - def _link_managed_parent(self, info: WorkflowTraceInfo, tenant_id: str) -> None: - """A legacy child cannot await a parent exported by a different pipeline.""" + def link_managed_parent(self, info: WorkflowTraceInfo) -> WorkflowTraceInfo: + """A provider-exported nested workflow names the Agent-collected parent it cannot attach to.""" + tenant_id = info.tenant_id parent_run_id, parent_node_id = info.resolved_parent_context if not parent_run_id: - return + return info with self.repository.sessions() as session: parent = session.scalar( select(ApplicationTraceRun) @@ -334,12 +302,12 @@ def _link_managed_parent(self, info: WorkflowTraceInfo, tenant_id: str) -> None: ) ) if parent is None: - return + return info 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 metadata = dict(info.metadata) metadata.pop("parent_trace_context", None) metadata["linked_parent_workflow_run_id"] = parent_run_id @@ -347,18 +315,23 @@ 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.""" + """Take records of Agent-managed executions; ``False`` leaves the record to its provider.""" 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) + # 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 = info.message_data + owner_app_id = message.app_id if message is not None else info.metadata.get("app_id") + app = self._app(owner_app_id) if isinstance(owner_app_id, str) and owner_app_id else None + return app is not None and self._managed_owner(app, message) with self.repository.sessions() as session: app = session.scalar(select(App).where(App.id == app_id)) if app is None: @@ -375,8 +348,6 @@ 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, @@ -384,10 +355,9 @@ def capture_terminal(self, info: BaseTraceInfo) -> bool: 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 @@ -397,23 +367,21 @@ def capture_terminal(self, info: BaseTraceInfo) -> bool: app.mode == AppMode.AGENT and message is not None and info.metadata.get("message_execution_state") not in ("succeeded", "failed", "cancelled") + and agent_stop_requested(message) ): - 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") + synthetic = frozenset({span_id_for(f"{source_id}:llm")}) if source_kind == "message" else frozenset() trace = remap_outer_trace( - trace, trace_id=run.trace_id, capture_content=run.capture_content, agent_message=source_kind == "message" + trace, trace_id=run.trace_id, capture_content=run.capture_content, synthetic_span_ids=synthetic ) 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.repository.terminal(run.id, trace.model_dump(mode="json"), trace.root.span_id) self.schedule(run.id) return True @@ -426,63 +394,76 @@ def settle_preview(self, run: ApplicationTraceRun) -> None: 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, + span = Span( + span_id=run.root_span_id, + parent_span_id=None, name=node.title or "Agent preview", - kind=CanonicalSpanKind.AGENT, + kind=SpanKind.AGENT, start_time=started_at, end_time=ended_at, + status=SpanStatus.ERROR if error else SpanStatus.OK, + error=error if run.capture_content else "Node execution failed" if error else None, 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"}, + attributes={"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,) + trace = Trace( + trace_id=run.trace_id, + external_trace_id=None, + session_id=run.session_id, + user_id=None, + parent=None, + spans=(span,), ) self.repository.terminal(run.id, trace.model_dump(mode="json"), run.root_span_id) - def assemble(self, run: ApplicationTraceRun) -> CanonicalTrace: + def assemble(self, run: ApplicationTraceRun) -> Trace: if run.outer_trace is None: raise TraceCollectionError("tracing_terminal_missing") - outer = CanonicalTrace.model_validate(run.outer_trace) + # Runs captured before the upgrade stored the previous trace shape, named by its root span. + outer = ( + legacy_outer_trace(run.outer_trace) + if "root_span_id" in run.outer_trace + else Trace.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} + internal_by_invocation: dict[str, dict[str, Span]] = {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) + current = internal_by_invocation[batch.invocation_id].get(span.span_id) if current is not None and current != span: raise TraceCollectionError("tracing_span_conflict") - internal_by_invocation[batch.invocation_id][span.id] = span + internal_by_invocation[batch.invocation_id][span.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): + 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 TraceCollectionError("tracing_invocation_parent_invalid") - root_children = [child for child in children if child.parent_id == invocation.span_id] + root_children = [child for child in children if child.parent_span_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] + errors = [child.error for child in root_children if child.status is SpanStatus.ERROR] spans.append( - CanonicalSpan( - id=invocation.span_id, - parent_id=invocation.parent_span_id, + Span( + span_id=invocation.span_id, + parent_span_id=invocation.parent_span_id, name=f"Agent attempt {invocation.attempt + 1}", - kind=CanonicalSpanKind.AGENT, + kind=SpanKind.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, + 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, - metadata={ + attributes={ "agent_id": invocation.agent_id, "attempt": invocation.attempt, "node_execution_id": invocation.node_execution_id, @@ -491,7 +472,7 @@ def assemble(self, run: ApplicationTraceRun) -> CanonicalTrace: ) ) spans.extend(children) - trusted = { + trusted: dict[str, Any] = { "tenant_id": run.tenant_id, "app_id": run.app_id, "session_id": run.session_id, @@ -500,12 +481,14 @@ def assemble(self, run: ApplicationTraceRun) -> CanonicalTrace: "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( + spans = [span.model_copy(update={"attributes": {**span.attributes, **trusted}}) for span in spans] + return Trace( trace_id=run.trace_id, + external_trace_id=None, session_id=run.session_id, - root_span_id=outer.root_span_id, - spans=order_spans(outer.root_span_id, spans), + user_id=outer.user_id, + parent=None, + spans=order_spans(outer.root.span_id, spans), ) @contextmanager @@ -553,8 +536,6 @@ 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: @@ -572,16 +553,15 @@ def deliver(self, run_id: str) -> None: try: try: trace = self.assemble(run) - except ValueError as error: + 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: + provider = self._frozen_provider(run) + result = provider.send(decode_config(run.tenant_id, type(provider), run.provider_config), trace) + except ExportOutcomeUnknownError: self.repository.complete(run_id, token, state="uncertain", error_code="provider_acceptance_unknown") - except RetryableTraceDispatchError: + except ExportRetryableError: if run.attempts >= dify_config.OPS_AGENT_TRACE_MAX_ATTEMPTS: self.repository.complete(run_id, token, state="failed", error_code="provider_retry_exhausted") else: @@ -594,6 +574,10 @@ def deliver(self, run_id: str) -> None: ) except TraceCollectionError as error: self.repository.complete(run_id, token, state="failed", error_code=error.code) + except ProviderConfigUnreadableError: + self.repository.complete(run_id, token, state="failed", error_code="tracing_config_unreadable") + except ExportRejectedError: + self.repository.complete(run_id, token, state="failed", error_code="provider_export_rejected") except Exception as error: logger.warning( "Application trace export failed: run_id=%s provider=%s error_type=%s", @@ -605,6 +589,27 @@ def deliver(self, run_id: str) -> None: else: self.repository.complete(run_id, token, state="accepted", result=result) + def _frozen_provider(self, run: ApplicationTraceRun) -> TraceProvider[Any]: + try: + return provider_class(run.provider)(self._egress) + except (UnknownTraceProviderError, TraceProviderUnavailableError): + raise TraceCollectionError("tracing_provider_unavailable") from None + + @staticmethod + def _frozen_capabilities(provider_id: str) -> dict[str, bool | str] | None: + try: + return capabilities_payload(provider_id) + except (AppTracingConfigInvalidProviderError, AppTracingConfigProviderUnavailableError): + return None + + def _frozen_project_url(self, run: ApplicationTraceRun) -> str | None: + """Present the run's own destination, never the app's current one.""" + try: + provider = provider_class(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]: @@ -664,7 +669,7 @@ def _status_for_run( "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), + "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, @@ -673,7 +678,7 @@ def _status_for_run( "error_code": run.error_code, "accepted_at": run.accepted_at, "queryable": None, - "capabilities": provider_capabilities(run.provider) if run.provider else None, + "capabilities": self._frozen_capabilities(run.provider) if run.provider else None, } def retry( @@ -702,4 +707,10 @@ def retry( def application_tracing_service() -> ApplicationTracingService: - return ApplicationTracingService(ApplicationTraceRepository(sessionmaker(bind=db.engine, expire_on_commit=False))) + sessions = sessionmaker(bind=db.engine, expire_on_commit=False) + return ApplicationTracingService( + ApplicationTraceRepository(sessions), + assembler=TraceAssembler(sessions, files_url=dify_config.FILES_URL, enterprise_details=False), + builder=CanonicalTraceBuilder(), + egress=SSRFEgress(), + ) diff --git a/api/services/workflow_persistence.py b/api/services/workflow_persistence.py index 28bb6750b83d0f..1cf938839b8a4f 100644 --- a/api/services/workflow_persistence.py +++ b/api/services/workflow_persistence.py @@ -9,12 +9,11 @@ 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 +21,7 @@ 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 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/tasks/ops_trace_task.py b/api/tasks/ops_trace_task.py index 1295b3f97b4ba3..8dbfd4e38da9e2 100644 --- a/api/tasks/ops_trace_task.py +++ b/api/tasks/ops_trace_task.py @@ -1,157 +1,29 @@ -""" -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; +``sweep_trace_payloads`` removes any payload this task never sees. """ -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..e8818a1b212777 --- /dev/null +++ b/api/tasks/trace_delivery_task.py @@ -0,0 +1,112 @@ +"""Celery entry points of the trace pipeline, on the ``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_storage import storage +from services.tracing.assembly import SourceNotFoundError, SourceNotReadyError +from services.tracing.delivery import ExportJob, TraceCollectorUnavailableError, trace_delivery_service +from services.tracing.emitter import PAYLOAD_PREFIX + +logger = logging.getLogger(__name__) + +REDELIVERY_LIMIT = 10 +PAYLOAD_RETENTION = timedelta(days=2) +_LEGACY_PAYLOAD_PREFIX = "ops_trace/" +_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(queue="ops_trace", 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: + trace_delivery_service().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(queue="ops_trace") +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: + trace_delivery_service().retry(job) + finally: + _delete_quietly(job_key) + + +@shared_task(queue="ops_trace") +def sweep_trace_payloads() -> None: + """Delete payloads a crashed delivery left behind, and payloads of the previous pipeline.""" + cutoff = (datetime.now(UTC) - PAYLOAD_RETENTION).date() + # Directory listings need the trailing slash on some storage backends. + for prefix in (f"{PAYLOAD_PREFIX}/", _LEGACY_PAYLOAD_PREFIX): + try: + keys = storage.scan(prefix, files=True) + except (FileNotFoundError, NotImplementedError): + continue + for key in keys: + if not key.startswith(prefix): + continue + if prefix == _LEGACY_PAYLOAD_PREFIX or _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/services/test_message_suggested_questions_failures.py b/api/tests/e2e_tests/containers/services/test_message_suggested_questions_failures.py index ac9b2203562ef5..62c9d6b3ba9f2c 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 @@ -59,7 +58,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 +138,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 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..a56c4414d7988a 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 @@ -57,6 +57,7 @@ MessageSuggestedQuestionsService, ) from services.recommended_app_package_service import RecommendedAppPackageService +from services.tracing.provider_configs import ProviderConfigGateway 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 @@ -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=ProviderConfigGateway(RecordingEgress([], grpc_target=None)), ), oauth=_build_oauth_server_service(database_client=sqlite_session_factory, redis=redis), recommended_packages=RecommendedAppPackageService( 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/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..35bd1438ddb153 100644 --- a/api/tests/unit_tests/configs/test_app_config.py +++ b/api/tests/unit_tests/configs/test_app_config.py @@ -15,32 +15,9 @@ from enums import DeploymentEdition -def test_ops_trace_config_rejects_parent_context_ttl_shorter_than_retry_window() -> None: - with pytest.raises(ValidationError, match="must cover the retry window"): - OpsTraceConfig( - OPS_TRACE_UNIFIED_ENABLED=True, - OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=4, - OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS=5, - OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=19, - ) - - -def test_ops_trace_config_skips_parent_context_validation_when_unified_tracing_is_disabled() -> None: - OpsTraceConfig( - OPS_TRACE_UNIFIED_ENABLED=False, - OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=4, - OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS=5, - OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=19, - ) - - -def test_ops_trace_config_accepts_parent_context_ttl_covering_retry_window() -> None: - OpsTraceConfig( - OPS_TRACE_UNIFIED_ENABLED=True, - OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=4, - OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS=5, - OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=20, - ) +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) @pytest.mark.parametrize( @@ -135,8 +112,8 @@ 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.ENABLE_TRACE_PAYLOAD_SWEEP_TASK is True assert config.GRAPH_ENGINE_MAX_WORKERS == 10 # annotated field with custom configured value @@ -149,6 +126,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..0158c381f3e137 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 @@ -24,6 +25,7 @@ from extensions import ext_redis from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.base import TypeBase +from services.tracing.provider_configs import ProviderConfigGateway from tests.helpers.unit.config_override import apply_config_overrides from tests.helpers.unit.rbac_fakes import RBACDomain, build_rbac_domain @@ -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=ProviderConfigGateway(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..ee9adf85ebd052 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,25 @@ -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 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 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.""" @@ -214,8 +197,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 +286,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..785600efd13fba 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 @@ -8,6 +8,7 @@ from uuid import UUID import pytest +from dify_trace_sdk.testing import RecordingEgress from flask import Flask from flask_restx import Resource from pydantic import ValidationError @@ -24,6 +25,8 @@ 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.assembly import TraceAssembler +from services.tracing.canonical import CanonicalTraceBuilder from services.tracing.service import ApplicationTracingService from tests.helpers.unit.config_override import apply_config_overrides @@ -54,7 +57,12 @@ def services( account_application_services: ApplicationServices, sqlite_session_factory: sessionmaker[Session], ) -> RunCase: - runs = ApplicationTracingService(ApplicationTraceRepository(sqlite_session_factory)) + runs = ApplicationTracingService( + ApplicationTraceRepository(sqlite_session_factory), + assembler=TraceAssembler(sqlite_session_factory, files_url="https://files.test", enterprise_details=False), + builder=CanonicalTraceBuilder(), + egress=RecordingEgress([], grpc_target=None), + ) case = RunCase(account_application_services, runs) def resolve(context: RequestContext, agent_id: str) -> str: 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..a43e8b16c434d4 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, @@ -44,7 +50,9 @@ AppTracingConfigService, AppTracingConfigVerificationFailedError, ) +from services.tracing.provider_configs import ProviderConfigGateway from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_app 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=ProviderConfigGateway(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..9c998fad151a43 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 @@ -4,6 +4,7 @@ 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 @@ -11,12 +12,25 @@ from controllers.inner_api import bp from models.application_trace import ApplicationTraceBatch +from repositories.application_trace_repository import ApplicationTraceRepository +from services.tracing.assembly import TraceAssembler +from services.tracing.canonical import CanonicalTraceBuilder 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 +def _service(sessions: sessionmaker[Session], repository: ApplicationTraceRepository) -> ApplicationTracingService: + return ApplicationTracingService( + repository, + assembler=TraceAssembler(sessions, files_url="https://files.test", enterprise_details=False), + builder=CanonicalTraceBuilder(), + # Ingesting batches never exports; an unexpected request fails the test. + egress=RecordingEgress([], grpc_target=None), + ) + + @pytest.fixture def client() -> Iterator[FlaskClient]: app = Flask(__name__) @@ -36,7 +50,7 @@ def test_trace_http_ack_follows_commit_and_replay_is_idempotent( 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) + service = _service(sqlite_session_factory, collection.repository) monkeypatch.setattr("controllers.inner_api.agent.tracing.application_tracing_service", lambda: service) scheduled = [] @@ -73,7 +87,7 @@ def test_trace_http_rejects_invalid_span_and_unknown_owner( collection = seed(sqlite_session_factory, "http-validation") invocation = collection.invocation() payload = batch_for(collection.run, invocation).model_dump(mode="json") - service = ApplicationTracingService(collection.repository) + service = _service(sqlite_session_factory, collection.repository) monkeypatch.setattr("controllers.inner_api.agent.tracing.application_tracing_service", lambda: service) payload["spans"][0]["end_time_ns"] = 0 response = client.post("/inner/api/agent/tracing/batches", json=payload, headers={"X-Inner-Api-Key": "inner"}) 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..9f2f0bafc60747 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,11 @@ ) 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, 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 +79,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 +324,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 +338,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 +350,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 +364,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 +391,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 +422,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 +568,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 +696,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 +807,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 +856,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 +907,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 +973,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 +1061,7 @@ def prepare(self): stream=False, invoke_from=InvokeFrom.DEBUGGER, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) @@ -1078,7 +1117,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 +1167,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 +1225,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 +1295,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 +1357,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 +1408,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 +1424,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 +1438,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 +1456,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 +1485,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 +1513,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 +1661,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 +1712,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 +1756,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 +1955,40 @@ 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", + 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 +2038,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 +2081,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..7d734b9af7b6c3 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,8 +16,10 @@ 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 +from sqlalchemy.orm import sessionmaker from clients.agent_backend import ( DIFY_CONFIG_LAYER_ID, @@ -34,13 +36,14 @@ 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.emitter import TraceScopes 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 @@ -53,6 +56,14 @@ class _LayerDeps: layers: dict[str, dict[str, object]] +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + @pytest.fixture(autouse=True) def _application_tracing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) @@ -155,9 +166,12 @@ def prepare(value: AgentAppRuntimeBuildContext) -> ApplicationTracingContext: return tracing monkeypatch.setattr(RuntimeApplicationTracingHooks, "prepare_agent_app", staticmethod(prepare)) - hooks = RuntimeApplicationTracingHooks() app = Flask(__name__) - bind_application_tracing_hooks(app, hooks) + bind_tracing( + app, + scopes=TraceScopes(emitter=_TraceRecorder(), session_factory=sessionmaker(), enterprise_enabled=lambda: False), + hooks=RuntimeApplicationTracingHooks(), + ) with app.app_context(): result = AgentAppRuntimeRequestBuilder(dify_tools_builder=WorkflowAgentDifyToolsBuilder()).build(context) 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 e4cc1f3d758ae1..2dfb989a6367a9 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())] @@ -5349,7 +5450,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, @@ -5394,7 +5494,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), ): @@ -5428,7 +5527,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() @@ -5439,7 +5538,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: @@ -5561,7 +5660,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) @@ -5592,10 +5691,10 @@ def test_on_retrieval_end_retries_postgres_deadlock(self, retrieval: DatasetRetr patch.object(retrieval, "_send_trace_task") as mock_trace, ): mock_sessionmaker.return_value.begin.side_effect = [deadlock, nullcontext(retry_session)] - 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()) assert mock_sessionmaker.return_value.begin.call_count == 2 - 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_runtime_request_builder.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py index a9471bfe52346d..79fea3781db32d 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 @@ -11,6 +11,7 @@ from dify_agent.protocol import DIFY_AGENT_HISTORY_LAYER_ID, DIFY_AGENT_MODEL_LAYER_ID, DIFY_AGENT_OUTPUT_LAYER_ID from dify_agent.protocol.snapshot import SessionSnapshot from flask import Flask +from sqlalchemy.orm import sessionmaker from clients.agent_backend import ( DIFY_CONFIG_LAYER_ID, @@ -20,7 +21,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 +30,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,6 +51,7 @@ DeclaredOutputType, WorkflowNodeJobConfig, ) +from services.tracing.emitter import CeleryTraceEmitter, TraceScopes 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 @@ -272,7 +274,13 @@ def prepare(value: WorkflowAgentRuntimeBuildContext, conversation_id: str | None monkeypatch.setattr(RuntimeApplicationTracingHooks, "prepare_workflow_agent", staticmethod(prepare)) hooks = RuntimeApplicationTracingHooks() app = Flask(__name__) - bind_application_tracing_hooks(app, hooks) + bind_tracing( + app, + scopes=TraceScopes( + emitter=CeleryTraceEmitter(), session_factory=sessionmaker(), enterprise_enabled=lambda: False + ), + hooks=hooks, + ) with app.app_context(): result = WorkflowAgentRuntimeRequestBuilder(dify_tools_builder=WorkflowAgentDifyToolsBuilder()).build(context) 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 842802148132ee..e91e22bc8e76f7 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 @@ -1753,6 +1753,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/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..a62cf0e3eb8dd9 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,8 @@ WorkflowNodeExecutionStatus, WorkflowNodeExecutionTriggeredFrom, ) +from services.tracing.assembly import TraceAssembler +from services.tracing.records import DraftNodeExecutionTrace # --------------------------------------------------------------------------- # Helpers @@ -235,8 +247,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 +259,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 +267,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 +285,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 +300,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 +315,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 → TraceAssembler → 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 +404,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 +421,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 +508,30 @@ 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 = TraceAssembler( + 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 +540,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..763908953eab05 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,15 @@ 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 models.workflow import WorkflowRun +from services.tracing.records import ( DatasetRetrievalTraceInfo, DraftNodeExecutionTrace, GenerateNameTraceInfo, @@ -22,13 +30,7 @@ WorkflowNodeTraceInfo, WorkflowTraceInfo, ) -from enterprise.telemetry.enterprise_trace import EnterpriseOtelTrace -from enterprise.telemetry.entities import ( - EnterpriseTelemetryCounter, - EnterpriseTelemetryEvent, - EnterpriseTelemetryHistogram, - EnterpriseTelemetrySpan, -) +from tests.helpers.unit.model_factories import make_message # --------------------------------------------------------------------------- # Fixtures @@ -59,8 +61,21 @@ def trace_handler(mock_exporter): _T1 = datetime(2024, 1, 10, 12, 0, 5, tzinfo=UTC) +def _workflow_run() -> WorkflowRun: + run = WorkflowRun(id="run-001", finished_at=_T1) + run.created_at = _T0 + return run + + def make_workflow_info(**overrides) -> WorkflowTraceInfo: defaults: dict[str, Any] = { + "workflow_data": _workflow_run(), + "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 +156,12 @@ def make_draft_node_info(**overrides) -> DraftNodeExecutionTrace: def make_message_info(**overrides) -> MessageTraceInfo: defaults: dict[str, Any] = { "message_id": "msg-001", + "message_data": make_message(message_id="msg-001", app_id="app-001"), + "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 +193,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 +221,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 +252,7 @@ def make_dataset_retrieval_info(**overrides) -> DatasetRetrievalTraceInfo: } ], "inputs": "search query", + "error": None, "metadata": { "app_id": "app-001", "tenant_id": "tenant-abc", @@ -266,6 +296,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) 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 ab18ed81c9c089..fdff1c556983b9 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -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 @@ -156,6 +155,9 @@ from services.tag_application_service import TagApplicationService from services.tool_file_download_service import ToolFileDownloadService from services.tracing import runtime_bridge +from services.tracing.egress import SSRFEgress +from services.tracing.emitter import CeleryTraceEmitter +from services.tracing.provider_configs import ProviderConfigGateway from services.upload_file_delivery_service import UploadFileDeliveryService from services.web_authentication_adapters import ( AccountWebAuthenticationSecurityGateway, @@ -170,6 +172,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 @@ -311,6 +314,40 @@ def construct_trace_service() -> Never: assert construction_calls == [] +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( sqlite_session_factory: sessionmaker[Session], ) -> None: @@ -639,7 +676,9 @@ 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) + provider = services.app_tracing_configs._provider + assert isinstance(provider, ProviderConfigGateway) + assert isinstance(provider._egress, SSRFEgress) def test_build_application_services_wires_workflow_app_log_boundary( 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..78b68c29f881c8 --- /dev/null +++ b/api/tests/unit_tests/extensions/test_tracing.py @@ -0,0 +1,96 @@ +from collections.abc import Mapping + +import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceEvent +from flask import Flask + +from extensions.tracing import ( + application_tracing_hooks, + bind_tracing, + create_trace_scope, + create_trace_scope_for_app, +) +from services.tracing.emitter import correlation_from_extras +from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks +from tests.helpers.unit.config_override import apply_config_overrides + + +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +class _Scopes: + def __init__(self, *, fails: bool) -> None: + self.fails = fails + self.policies: list[str | None] = [] + + def __call__( + self, + *, + tenant_id: str, + app_id: str | None, + app_tracing: str | None, + user_id: str | None, + extras: Mapping[str, object], + ) -> TraceScope | None: + self.policies.append(app_tracing) + if self.fails: + raise RuntimeError("policy lookup failed") + return TraceScope( + emitter=_Recorder(), + tenant_id=tenant_id, + app_id=app_id, + user_id=user_id, + correlation=correlation_from_extras(extras), + ) + + def for_app( + self, *, tenant_id: str, app_id: str | None, user_id: str | None, extras: Mapping[str, object] + ) -> TraceScope | None: + return self(tenant_id=tenant_id, app_id=app_id, app_tracing=None, user_id=user_id, extras=extras) + + +def _app(*, fails: bool) -> Flask: + app = Flask(__name__) + bind_tracing(app, scopes=_Scopes(fails=fails), hooks=RuntimeApplicationTracingHooks()) + return app + + +def test_scopes_come_from_the_bound_factory() -> None: + with _app(fails=False).app_context(): + scope = create_trace_scope(tenant_id="tenant-1", app_id="app-1", app_tracing=None, user_id="u", extras={}) + 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 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(fails=True).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_follow_the_agent_tracing_switch(monkeypatch: pytest.MonkeyPatch) -> None: + app = _app(fails=False) + + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) + with app.app_context(): + assert application_tracing_hooks() is None + + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + with app.app_context(): + assert isinstance(application_tracing_hooks(), RuntimeApplicationTracingHooks) + 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/test_application_trace_repository.py b/api/tests/unit_tests/repositories/test_application_trace_repository.py index d9eb6469526faa..06f3bcf80eece2 100644 --- a/api/tests/unit_tests/repositories/test_application_trace_repository.py +++ b/api/tests/unit_tests/repositories/test_application_trace_repository.py @@ -1,5 +1,5 @@ import json -from collections.abc import Iterator +from collections.abc import Callable, Iterator from dataclasses import dataclass from datetime import timedelta from typing import Literal @@ -7,18 +7,22 @@ import pytest from aliyun.log.logclient import LogClient +from dify_trace_sdk.events import WorkflowRunCompleted +from dify_trace_sdk.provider import SendResult 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 extensions.logstore.aliyun_logstore import AliyunLogStore 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.model import App, AppMode, Message from models.workflow import WorkflowNodeExecutionModel, WorkflowRun +from repositories import application_trace_repository as repository_module from repositories.application_trace_repository import TraceCollectionError -from services.tracing.provider_dispatch import TraceExportResult +from services.tracing.records import ModerationTraceInfo from services.tracing.recovery import _workflow_task from services.tracing.service import ApplicationTracingService from tests.helpers.unit.config_override import apply_config_overrides @@ -28,8 +32,12 @@ 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_no_egress as no_egress 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 ( + _application_trace_collection_tracing_service as tracing_service, +) from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection """Configured LogStore readers work without SQL mirrors and retain source ownership.""" @@ -54,7 +62,7 @@ def start(self) -> ApplicationTraceRun: return run def service(self) -> ApplicationTracingService: - return ApplicationTracingService(self.case.repository) + return tracing_service(self.case, no_egress()) def status(self, node_execution_id: str | None = None) -> dict[str, object]: return self.service().status( @@ -280,7 +288,9 @@ def test_workflow_node_deep_link_is_scoped_to_durable_invocation(stored: StoredS run.id, lease[1], state="accepted", - result=TraceExportResult("provider-id", span_urls={invocation.span_id: "https://trace.test/span"}), + result=SendResult( + trace_id="provider-id", trace_url=None, 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"): @@ -290,10 +300,14 @@ def test_workflow_node_deep_link_is_scoped_to_durable_invocation(stored: StoredS 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 + event = _workflow_task(stored.case.repository, run, naive_utc_now()) + assert event is not None + assert event.payload == WorkflowRunCompleted( + workflow_run_id=run.source_id, conversation_id=None, total_tokens=None + ) + assert (event.tenant_id, event.app_id) == (run.tenant_id, run.app_id) + assert event.user_id == stored.row["created_by"] + assert event.correlation.session_id == run.session_id stored.row["status"] = "running" assert _workflow_task(stored.case.repository, run, naive_utc_now()) is None else: @@ -305,12 +319,12 @@ 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)) + info = workflow_info(_dict_to_workflow_run(stored.row), None) assert stored.service().capture_terminal(info) is True result = stored.status() assert result["status"] == "failed" @@ -322,7 +336,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 @@ -340,15 +354,21 @@ def test_unindexed_workflow_uses_owned_definition_for_safe_legacy_decision( @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=message, + metadata={"message_id": message.id}, flagged=False, action="direct_output", preset_response="", @@ -374,33 +394,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(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( @@ -434,7 +436,92 @@ 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 check, 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(repository_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) -> ApplicationTraceRun | None: + return case.repository.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, + ) + + 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[ApplicationTraceRun | 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 == "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) + winners: list[ApplicationTraceRun | None] = [] + race_before_insert(monkeypatch, lambda: winners.append(start_message(case, "started-session"))) + current = 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 + ) + assert current is not None + assert winners[0] is not None + assert current.id == winners[0].id + assert current.state == "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(repository_module, "uuidv7", lambda: UUID(other.run.id)) + with pytest.raises(IntegrityError): + start_message(case, "colliding-session") + 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( 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..2f7078091617a0 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 services.tracing.provider_configs import ProviderConfigGateway 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=ProviderConfigGateway(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..0df12fc5f4596d 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,9 @@ AppTracingConfigRecord, AppTracingConfigService, AppTracingConfigVerificationFailedError, + PresentedTracingConfig, ) +from services.tracing.provider_configs import ProviderConfigGateway """Public tracing operations keep credentials and policy in one transaction.""" @@ -33,7 +35,7 @@ class Case: sessions: sessionmaker[Session] store: SQLAlchemyAppTracingConfigRepository service: AppTracingConfigService - gateway: OpsTraceManagerGateway + gateway: ProviderConfigGateway reader: ConsoleAppRepository def policy(self) -> AppTraceSettings: @@ -61,7 +63,7 @@ def case(sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.Monk ), ) ) - gateway = OpsTraceManagerGateway() + gateway = ProviderConfigGateway(RecordingEgress([], grpc_target=None)) def prepare(*, tracing_config: dict[str, object], **_: object) -> dict[str, object]: return dict(tracing_config) @@ -257,12 +259,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..43413d13bd4f66 100644 --- a/api/tests/unit_tests/services/test_conversation_name_generator.py +++ b/api/tests/unit_tests/services/test_conversation_name_generator.py @@ -121,12 +121,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" 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..a3f3ede270cde4 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 @@ -25,6 +25,7 @@ from core.plugin.impl.base import _get_plugin_daemon_request_timeout from extensions.ext_database import db from extensions.ext_redis import RedisClientWrapper +from extensions.tracing import bind_tracing from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent from graphon.model_runtime.entities.model_entities import ModelType from models.provider import Provider, TenantDefaultModel @@ -35,6 +36,8 @@ _default_suggested_questions_model_parameters, ) from services.message_suggested_questions_service import SuggestedQuestionsContext +from services.tracing.emitter import TraceScopes +from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks from tests.helpers.unit.core.model_fixtures import make_model_instance from tests.helpers.unit.model_factories import make_app @@ -196,11 +199,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 +247,19 @@ def record_session(session: Session, _transaction: SessionTransaction, _connecti sessions.append(session) factory = db.session.session_factory + recorder = _TraceRecorder() + bind_tracing( + current_app, + scopes=TraceScopes(emitter=recorder, session_factory=sqlite_session_factory, enterprise_enabled=lambda: False), + hooks=RuntimeApplicationTracingHooks(), + ) 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 = SuggestedQuestionsGenerator().generate( PreparedHistory(prompts=()), history_model=make_model_instance(provider="invalid/provider", model="unused"), context=context, @@ -242,7 +268,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] diff --git a/api/tests/unit_tests/services/test_workflow_persistence.py b/api/tests/unit_tests/services/test_workflow_persistence.py index cf508c3f14e86a..fa6aa73d73d59a 100644 --- a/api/tests/unit_tests/services/test_workflow_persistence.py +++ b/api/tests/unit_tests/services/test_workflow_persistence.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, WorkflowRunCompleted from flask import Flask from sqlalchemy import event as sqlalchemy_event from sqlalchemy import select @@ -14,8 +16,6 @@ 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 @@ -25,6 +25,7 @@ ) from core.workflow.system_variables import SystemVariableKey, build_system_variables 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,6 +73,7 @@ DifyAPISQLAlchemyWorkflowNodeExecutionRepository, ) from services import workflow_persistence as persistence_mod +from services.tracing.emitter import TraceScopes from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks from services.workflow_persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer from tests.helpers.unit.config_override import apply_config_overrides @@ -138,6 +141,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 +277,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 +309,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 +333,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 +356,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 {}) @@ -903,9 +930,16 @@ def freeze( ) monkeypatch.setattr(RuntimeApplicationTracingHooks, "freeze_workflow_run", staticmethod(freeze)) - hooks = RuntimeApplicationTracingHooks() app = Flask(__name__) - bind_application_tracing_hooks(app, hooks) + bind_tracing( + app, + scopes=TraceScopes( + emitter=_TraceRecorder(workflow_repo), + session_factory=sessionmaker(), + enterprise_enabled=lambda: False, + ), + hooks=RuntimeApplicationTracingHooks(), + ) with app.app_context(): layer._handle_graph_run_started() @@ -1033,11 +1067,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 +1091,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] - trace_task.execute() + 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()) - 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", - } + layer.on_event(GraphRunPausedEvent(reasons=[SchedulingPause(message="pause")])) + + 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_assembly.py b/api/tests/unit_tests/services/tracing/test_assembly.py new file mode 100644 index 00000000000000..4c290f582d7b33 --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_assembly.py @@ -0,0 +1,205 @@ +from datetime import UTC, datetime, timedelta +from decimal import Decimal + +import pytest +from dify_trace_sdk.events import ( + MessageCompleted, + ModerationChecked, + ToolInvoked, + TraceCorrelation, + TraceEvent, + TracePayload, +) +from sqlalchemy.orm import Session, sessionmaker + +from models.enums import ConversationFromSource, ConversationStatus, EndUserType, InvokeFrom +from models.model import App, AppMode, Conversation, EndUser, Message +from services.tracing.assembly import SourceNotFoundError, SourceNotReadyError, TraceAssembler +from services.tracing.records import MessageTraceInfo, ModerationTraceInfo, ToolTraceInfo + +TENANT_ID = "tenant-1" +APP_ID = "app-1" +MESSAGE_ID = "message-1" +CREATED_AT = datetime(2026, 1, 1, 8, 0, 0) + + +@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 _assembler(sessions: sessionmaker[Session]) -> TraceAssembler: + return TraceAssembler(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)) 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..6b872cce42b5bc --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_canonical.py @@ -0,0 +1,271 @@ +from datetime import UTC, datetime, timedelta +from decimal import Decimal + +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 models.model import Message +from models.workflow import WorkflowRun +from services.tracing.canonical import CanonicalTraceBuilder, message_trace_id +from services.tracing.records import ( + MessageTraceInfo, + PromptGenerationTraceInfo, + ToolTraceInfo, + 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() -> Message: + message = Message( + id=MESSAGE_ID, + app_id="app-1", + conversation_id=CONVERSATION_ID, + query="What is Dify?", + answer="An LLM app platform.", + model_provider="openai", + model_id="gpt-4o", + total_price=Decimal("0.0042"), + currency="USD", + error=None, + ) + message.created_at = STARTED_AT + return message + + +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_run() -> WorkflowRun: + run = WorkflowRun(id=WORKFLOW_RUN_ID, finished_at=STARTED_AT + timedelta(seconds=3)) + run.created_at = STARTED_AT + return run + + +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, + workflow_data=_workflow_run(), + 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_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 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..caf05c5b0ab3d2 --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_delivery.py @@ -0,0 +1,331 @@ +from collections.abc import Callable +from datetime import UTC, datetime, timedelta +from typing import override + +import pytest +from dify_trace_sdk.config import ProviderConfig +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 +from pydantic import ValidationError +from sqlalchemy.orm import sessionmaker + +from models.model import Message +from services.tracing.assembly import SourceNotReadyError +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.delivery import ( + Destination, + ExportJob, + TraceCollectorUnavailableError, + TraceDeliveryService, + retry_delay_seconds, +) +from services.tracing.records import BaseTraceInfo, MessageTraceInfo, WorkflowTraceInfo + +STARTED_AT = datetime(2026, 1, 1, 8, 0, 0) +MESSAGE_ID = "6f1c0b52-6ad4-4c4e-9d7c-2a3f0d0c7a11" + + +class _Config(ProviderConfig): + endpoint: str + + +class _Provider(TraceProvider[_Config]): + id = "scripted" + config_model = _Config + capabilities = ProviderCapabilities(native_cost=False, trace_url=False, span_url=False) + + def __init__(self, outcomes: list[Exception | None]) -> None: + super().__init__(RecordingEgress([], grpc_target=None)) + self.outcomes = outcomes + self.sent: list[Trace] = [] + + @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.sent.append(trace) + outcome = self.outcomes.pop(0) + if outcome is not None: + raise outcome + return SendResult(trace_id=trace.trace_id, trace_url=None, span_urls={}) + + +class _Assembler: + def __init__(self, result: BaseTraceInfo | Exception) -> None: + self.result = result + self.events: list[TraceEvent] = [] + + def assemble(self, event: TraceEvent) -> BaseTraceInfo: + self.events.append(event) + if isinstance(self.result, Exception): + raise self.result + return self.result + + +class _Collector: + def __init__(self, *, captures: bool, fails: bool) -> None: + self.captures = captures + self.fails = fails + self.seen: list[BaseTraceInfo] = [] + + def capture_terminal(self, info: BaseTraceInfo) -> bool: + if self.fails: + raise RuntimeError("ledger unavailable") + self.seen.append(info) + return self.captures + + def link_managed_parent(self, info: WorkflowTraceInfo) -> WorkflowTraceInfo: + return info + + +class _Service(TraceDeliveryService): + """Delivery with an in-memory destination and retry schedule.""" + + def __init__( + self, + *, + assembler: _Assembler, + destination: Destination | None, + enterprise_sink: Callable[[BaseTraceInfo], None] | None, + collector: _Collector | None, + ) -> None: + super().__init__( + session_factory=sessionmaker(), + assembler=assembler, + builder=CanonicalTraceBuilder(), + egress=RecordingEgress([], grpc_target=None), + enterprise_sink=enterprise_sink, + collector=collector, + max_export_attempts=3, + ) + self._destination = destination + self.scheduled: list[ExportJob] = [] + + @override + def destination(self, tenant_id: str, app_id: str) -> Destination | None: + return self._destination + + @override + def _schedule_retry(self, job: ExportJob) -> None: + self.scheduled.append(job) + + +def _event() -> TraceEvent: + return TraceEvent( + event_id="event-1", + tenant_id="tenant-1", + app_id="app-1", + 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 = None) -> MessageTraceInfo: + message = Message( + id=MESSAGE_ID, + app_id="app-1", + conversation_id="conversation-1", + answer="hello", + workflow_run_id=workflow_run_id, + ) + message.created_at = STARTED_AT + 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 _destination(outcomes: list[Exception | None]) -> Destination: + return Destination(provider=_Provider(outcomes), config=_Config(endpoint="https://collector.example")) + + +def test_every_sink_receives_the_record_and_the_provider_gets_the_trace() -> None: + destination = _destination([None]) + enterprise: list[BaseTraceInfo] = [] + collector = _Collector(captures=False, fails=False) + service = _Service( + assembler=_Assembler(_record()), destination=destination, enterprise_sink=enterprise.append, collector=collector + ) + + service.deliver(_event(), enterprise_delivered=False) + + assert len(enterprise) == 1 + assert len(collector.seen) == 1 + assert isinstance(destination.provider, _Provider) + (trace,) = destination.provider.sent + assert trace.root.name == "message" + assert service.scheduled == [] + + +def test_enterprise_sink_is_skipped_on_redelivery() -> None: + enterprise: list[BaseTraceInfo] = [] + service = _Service( + assembler=_Assembler(_record()), + destination=_destination([None]), + enterprise_sink=enterprise.append, + collector=None, + ) + + service.deliver(_event(), enterprise_delivered=True) + + assert enterprise == [] + + +def test_records_the_collector_owns_never_reach_the_provider() -> None: + destination = _destination([]) + service = _Service( + assembler=_Assembler(_record()), + destination=destination, + enterprise_sink=None, + collector=_Collector(captures=True, fails=False), + ) + + service.deliver(_event(), enterprise_delivered=False) + + assert isinstance(destination.provider, _Provider) + assert destination.provider.sent == [] + + +def test_an_unavailable_collector_asks_for_redelivery() -> None: + service = _Service( + assembler=_Assembler(_record()), + destination=_destination([]), + enterprise_sink=None, + collector=_Collector(captures=False, fails=True), + ) + + with pytest.raises(TraceCollectorUnavailableError): + service.deliver(_event(), enterprise_delivered=False) + + +def test_chatflow_message_is_covered_by_its_workflow_trace() -> None: + destination = _destination([]) + service = _Service( + assembler=_Assembler(_record(workflow_run_id="run-1")), + destination=destination, + enterprise_sink=None, + collector=None, + ) + + service.deliver(_event(), enterprise_delivered=False) + + assert isinstance(destination.provider, _Provider) + assert destination.provider.sent == [] + + +def test_source_not_ready_propagates_for_redelivery() -> None: + service = _Service( + assembler=_Assembler(SourceNotReadyError()), destination=_destination([]), enterprise_sink=None, collector=None + ) + + with pytest.raises(SourceNotReadyError): + service.deliver(_event(), enterprise_delivered=False) + + +def test_retryable_failure_schedules_the_next_attempt_until_the_limit() -> None: + service = _Service( + assembler=_Assembler(_record()), + destination=_destination([ExportRetryableError("busy"), ExportRetryableError("busy")]), + enterprise_sink=None, + collector=None, + ) + + service.deliver(_event(), enterprise_delivered=False) + + (job,) = service.scheduled + assert (job.attempt, job.provider_id, job.app_id) == (2, "scripted", "app-1") + exhausted = job.model_copy(update={"attempt": 3}) + destination = service.destination("tenant-1", "app-1") + assert destination is not None + service.export(exhausted, destination) + assert len(service.scheduled) == 1 + + +@pytest.mark.parametrize("failure", [ExportRejectedError("bad key"), ExportOutcomeUnknownError("timeout")]) +def test_rejected_or_unknown_outcomes_are_never_retried(failure: Exception) -> None: + service = _Service( + assembler=_Assembler(_record()), destination=_destination([failure]), enterprise_sink=None, collector=None + ) + + service.deliver(_event(), enterprise_delivered=False) + + assert service.scheduled == [] + + +def test_retry_is_dropped_when_the_app_now_traces_elsewhere() -> None: + trace = CanonicalTraceBuilder().build(_record()) + assert trace is not None + destination = _destination([]) + service = _Service(assembler=_Assembler(_record()), destination=destination, enterprise_sink=None, collector=None) + + service.retry(ExportJob(tenant_id="tenant-1", app_id="app-1", provider_id="langfuse", attempt=2, trace=trace)) + + assert isinstance(destination.provider, _Provider) + assert destination.provider.sent == [] + + +def test_export_job_round_trips_through_storage_format() -> None: + trace = CanonicalTraceBuilder().build(_record()) + assert trace is not None + job = ExportJob(tenant_id="tenant-1", app_id="app-1", provider_id="scripted", 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({}) + destination = _destination([]) + service = _Service( + assembler=_Assembler(invalid.value), destination=destination, enterprise_sink=None, collector=None + ) + + service.deliver(_event(), enterprise_delivered=False) + + assert isinstance(destination.provider, _Provider) + assert destination.provider.sent == [] + + +def test_unclassified_provider_failures_are_contained() -> None: + service = _Service( + assembler=_Assembler(_record()), + destination=_destination([RuntimeError("provider bug with secret-detail")]), + enterprise_sink=None, + collector=None, + ) + + service.deliver(_event(), enterprise_delivered=False) + + assert service.scheduled == [] 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_emitter.py b/api/tests/unit_tests/services/tracing/test_emitter.py new file mode 100644 index 00000000000000..cd9c464b8cb18d --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_emitter.py @@ -0,0 +1,202 @@ +from datetime import UTC, datetime + +import pytest +from dify_trace_sdk.events import ModerationChecked, TraceCorrelation, TraceEvent, WorkflowRunCompleted +from sqlalchemy.orm import sessionmaker + +from core.helper.trace_id_helper import ParentTraceContext +from services.tracing import emitter +from services.tracing.emitter import ( + DELIVER_TASK, + INLINE_EVENT_BYTES, + TRACE_QUEUE, + WORKFLOW_SETTLE_SECONDS, + CeleryTraceEmitter, + TraceScopes, + correlation_from_extras, + payload_key, + tracing_policy_enabled, +) + + +class _Broker: + def __init__(self) -> None: + self.sent: list[tuple[str, tuple[dict[str, str]], str]] = [] + self.countdowns: list[int] = [] + + def send_task(self, name: str, *, args: tuple[dict[str, str]], queue: str, countdown: int) -> None: + self.sent.append((name, args, queue)) + self.countdowns.append(countdown) + + +class _Storage: + def __init__(self) -> None: + self.files: dict[str, bytes] = {} + + def save(self, key: str, data: bytes) -> None: + self.files[key] = data + + +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _event(query: str) -> TraceEvent: + started_at = datetime(2026, 3, 4, 5, 6, tzinfo=UTC) + return TraceEvent( + event_id="event-1", + tenant_id="tenant-1", + app_id="app-1", + user_id=None, + occurred_at=started_at, + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + payload=ModerationChecked( + message_id="message-1", + inputs={}, + query=query, + flagged=False, + action="direct_output", + preset_response="", + started_at=started_at, + ended_at=started_at, + ), + ) + + +@pytest.mark.parametrize( + ("policy", "enabled"), + [ + ('{"enabled": true, "tracing_provider": "langfuse"}', True), + ('{"enabled": false, "tracing_provider": "langfuse"}', False), + ('{"enabled": true, "tracing_provider": null}', False), + ("not json", False), + (None, False), + ], +) +def test_tracing_policy(policy: str | None, enabled: bool) -> None: + assert tracing_policy_enabled(policy) is enabled + + +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_small_events_travel_inline(monkeypatch: pytest.MonkeyPatch) -> None: + broker, storage = _Broker(), _Storage() + monkeypatch.setattr(emitter, "celery_app", broker) + monkeypatch.setattr(emitter, "storage", storage) + event = _event("hello") + + CeleryTraceEmitter().emit(event) + + ((name, (envelope,), queue),) = broker.sent + assert (name, queue) == (DELIVER_TASK, TRACE_QUEUE) + assert TraceEvent.model_validate_json(envelope["event"]) == event + assert storage.files == {} + + +def test_large_events_travel_through_storage(monkeypatch: pytest.MonkeyPatch) -> None: + broker, storage = _Broker(), _Storage() + monkeypatch.setattr(emitter, "celery_app", broker) + monkeypatch.setattr(emitter, "storage", storage) + event = _event("x" * INLINE_EVENT_BYTES) + + CeleryTraceEmitter().emit(event) + + ((_, (envelope,), _),) = broker.sent + assert envelope == {"payload_key": payload_key(event.occurred_at, "events", "event-1")} + assert envelope["payload_key"] == "trace_payloads/2026-03-04/events/event-1.json" + assert TraceEvent.model_validate_json(storage.files[envelope["payload_key"]]) == event + + +def _scopes(*, enterprise: bool) -> tuple[TraceScopes, _Recorder]: + recorder = _Recorder() + return TraceScopes( + emitter=recorder, session_factory=sessionmaker(), enterprise_enabled=lambda: enterprise + ), recorder + + +def test_no_scope_exists_when_nothing_consumes_the_events() -> None: + scopes, _ = _scopes(enterprise=False) + + assert scopes(tenant_id="tenant-1", app_id="app-1", app_tracing=None, user_id=None, extras={}) is None + + +@pytest.mark.parametrize( + ("policy", "enterprise"), + [('{"enabled": true, "tracing_provider": "langfuse"}', False), (None, True)], +) +def test_scopes_stamp_owner_and_correlation(policy: str | None, enterprise: bool) -> None: + scopes, recorder = _scopes(enterprise=enterprise) + + scope = scopes( + tenant_id="tenant-1", + app_id="app-1", + app_tracing=policy, + user_id="end-user-session", + extras={"external_trace_id": "caller-trace"}, + ) + + 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.external_trace_id == "caller-trace" + assert scope.emitter is recorder + + +def test_app_scopes_without_an_app_depend_on_enterprise_telemetry() -> None: + scopes, _ = _scopes(enterprise=True) + + scope = scopes.for_app(tenant_id="tenant-1", app_id=None, user_id=None, extras={}) + + assert scope is not None + assert scope.app_id is None + + +def test_workflow_runs_wait_for_node_executions_to_land(monkeypatch: pytest.MonkeyPatch) -> None: + broker = _Broker() + monkeypatch.setattr(emitter, "celery_app", broker) + workflow = _event("hello").model_copy( + update={"payload": WorkflowRunCompleted(workflow_run_id="run-1", conversation_id=None, total_tokens=0)} + ) + + CeleryTraceEmitter().emit(_event("hello")) + CeleryTraceEmitter().emit(workflow) + + assert broker.countdowns == [0, WORKFLOW_SETTLE_SECONDS] + + +def test_inline_size_is_measured_in_bytes(monkeypatch: pytest.MonkeyPatch) -> None: + broker, storage = _Broker(), _Storage() + monkeypatch.setattr(emitter, "celery_app", broker) + monkeypatch.setattr(emitter, "storage", storage) + + CeleryTraceEmitter().emit(_event("\u4e2d" * (INLINE_EVENT_BYTES // 2))) + + ((_, (envelope,), _),) = broker.sent + assert "payload_key" in envelope diff --git a/api/tests/unit_tests/services/tracing/test_normalization.py b/api/tests/unit_tests/services/tracing/test_normalization.py index 870b812dc2cbc7..f3722350b6ceda 100644 --- a/api/tests/unit_tests/services/tracing/test_normalization.py +++ b/api/tests/unit_tests/services/tracing/test_normalization.py @@ -4,11 +4,17 @@ 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 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 +from services.tracing.normalization import ( + legacy_outer_trace, + normalize_agent_span, + order_spans, + remap_outer_trace, + safe_attributes, +) def agent_data( @@ -50,9 +56,10 @@ def test_content_off_is_safe_for_all_metadata_namespaces() -> 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 + 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: @@ -73,7 +80,9 @@ def test_json_encoded_tool_credentials_are_redacted_from_inputs_and_metadata() - 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]" + 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: @@ -98,6 +107,7 @@ def test_real_runtime_accounting_attributes_survive_content_off() -> None: "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, @@ -108,46 +118,79 @@ def test_real_runtime_accounting_attributes_survive_content_off() -> None: ), 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.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 -def span(identifier: str, parent: str | None, kind: CanonicalSpanKind, *, synthetic: bool = False) -> CanonicalSpan: - return CanonicalSpan( - id=identifier, - parent_id=parent, - name=identifier, +@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=CanonicalSpanStatus.OK, + status=SpanStatus.OK, inputs={"content": "private"}, outputs="private", - metadata={ - "prompt_tokens": 6, - "completion_tokens": 4, - "total_tokens": 10, - "total_price": "0.03", - "prompt": "private", - }, - synthetic=synthetic, + 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 test_two_native_model_calls_replace_synthetic_aggregate_generation() -> None: - outer = CanonicalTrace( - trace_id="old", +def trace_of(*spans: Span) -> Trace: + return Trace( + trace_id=span_id_for("old") * 2, + external_trace_id=None, session_id="session", - root_span_id="message", - spans=( - span("message", None, CanonicalSpanKind.CHAIN), - span("message:llm", "message", CanonicalSpanKind.LLM, synthetic=True), - ), + user_id="user", + parent=None, + spans=spans, ) - remapped = remap_outer_trace(outer, trace_id="1" * 32, capture_content=False, agent_message=True) + + +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( @@ -158,61 +201,163 @@ def test_two_native_model_calls_replace_synthetic_aggregate_generation() -> None "gen_ai.usage.output_tokens": 2, }, span_id=identifier, - parent_id=remapped.root_span_id, + 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] + 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(token_usage(item)["total"] for item in generations) == 10 - assert token_usage(combined[0]) == {} - assert combined[0].metadata["aggregate_total_tokens"] == 10 + 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 = 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) + 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 - assert remapped.spans[1].kind is CanonicalSpanKind.LLM - assert token_usage(remapped.spans[1])["total"] == 10 + 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, CanonicalSpanKind.CHAIN) + root = span("root", None, SpanKind.CHAIN) with pytest.raises(ValueError, match="parent is missing"): - order_spans("root", [root, span("orphan", "missing", CanonicalSpanKind.TOOL)]) + order_spans(root.span_id, [root, span("orphan", "missing", SpanKind.TOOL)]) with pytest.raises(ValueError, match="parent cycle"): - order_spans( - "root", [root, span("one", "two", CanonicalSpanKind.TOOL), span("two", "one", CanonicalSpanKind.TOOL)] - ) + 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", [root, root.model_copy(update={"name": "other"})]) + 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: - trace = CanonicalTrace( - trace_id="old", + 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", - 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" - ), + user_id=None, + parent=None, + spans=(root,), ) - 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" + 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/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..9a5ea8ed28df60 --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_provider_configs.py @@ -0,0 +1,272 @@ +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 dify_trace_sdk.egress import 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 import provider_configs, registry +from services.tracing.provider_configs import ( + ProviderConfigGateway, + ProviderConfigUnreadableError, + capabilities_payload, + configured_secrets, + decode_config, + encode_config, +) + +WORKSPACE = "tenant-1" +PUBLIC_KEY = "pk-lf-0123456789" +SECRET_KEY = "sk-lf-abcdefghijk" + + +def _encrypt(tenant_id: str, token: str) -> str: + return f"enc:{tenant_id}:{token}" + + +def _batch_decrypt(tenant_id: str, tokens: list[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") + return [token.removeprefix(prefix) for token in tokens] + + +@pytest.fixture(autouse=True) +def reversible_encryption(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(provider_configs, "encrypt_token", _encrypt) + monkeypatch.setattr(provider_configs, "batch_decrypt_token", _batch_decrypt) + + +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_stored_configs_encrypt_only_secrets_and_round_trip() -> None: + config = LangfuseConfig(public_key=PUBLIC_KEY, secret_key=SECRET_KEY, host="https://langfuse.example.com") + + stored = encode_config(WORKSPACE, config) + + assert stored["public_key"] == _encrypt(WORKSPACE, PUBLIC_KEY) + assert stored["host"] == "https://langfuse.example.com" + assert decode_config(WORKSPACE, LangfuseProvider, stored) == config + + +def test_a_config_encrypted_for_another_workspace_is_unreadable() -> None: + with pytest.raises(ProviderConfigUnreadableError): + decode_config("tenant-2", LangfuseProvider, _stored()) + + +def test_presentation_never_returns_secrets_or_derived_fields() -> None: + gateway = ProviderConfigGateway(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_new_configs_are_verified_and_store_derived_fields() -> None: + egress = RecordingEgress([_projects("project-9")], grpc_target=None) + gateway = ProviderConfigGateway(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", + }, + ) + + assert stored["project_key"] == "project-9" + 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 = ProviderConfigGateway(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 = ProviderConfigGateway(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 = ProviderConfigGateway(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): + ProviderConfigGateway(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 = ProviderConfigGateway(RecordingEgress([], grpc_target=None)) + + with pytest.raises(config_service.AppTracingConfigInvalidProviderError): + gateway.validate_provider("not-a-provider") + + +def test_capabilities_describe_the_provider_mapping() -> None: + payload = capabilities_payload("langfuse") + + assert payload["provider"] == "langfuse" + assert payload["native_cost"] is True + assert payload["export_ack"] == "accepted" + + +def test_empty_optional_secrets_are_stored_empty_and_not_reported_as_configured() -> None: + config = PhoenixConfig(api_key="", project="dify", endpoint="https://app.phoenix.arize.com") + + stored = encode_config(WORKSPACE, config) + + assert stored["api_key"] == "" + assert configured_secrets(PhoenixProvider, stored) == () + assert decode_config(WORKSPACE, PhoenixProvider, stored) == config + + +def test_uninstalled_providers_can_still_be_named_but_not_presented(monkeypatch: pytest.MonkeyPatch) -> None: + registry.provider_class.cache_clear() + monkeypatch.setattr(registry, "_entry_points", dict) + gateway = ProviderConfigGateway(RecordingEgress([], grpc_target=None)) + try: + gateway.validate_provider("langfuse") + with pytest.raises(config_service.AppTracingConfigProviderUnavailableError): + gateway.present_config(workspace_id=WORKSPACE, tracing_provider="langfuse", tracing_config=_stored()) + finally: + registry.provider_class.cache_clear() + + +def test_secrets_that_cannot_be_decrypted_make_the_config_unreadable(monkeypatch: pytest.MonkeyPatch) -> None: + class _MissingWorkspaceKeyError(Exception): + pass + + def fail(tenant_id: str, tokens: list[str]) -> list[str]: + raise _MissingWorkspaceKeyError(f"{tenant_id}: {len(tokens)} tokens") + + monkeypatch.setattr(provider_configs, "batch_decrypt_token", fail) + + with pytest.raises(ProviderConfigUnreadableError): + decode_config(WORKSPACE, LangfuseProvider, _stored()) + + +def test_configs_without_secrets_never_load_the_workspace_key(monkeypatch: pytest.MonkeyPatch) -> None: + def forbidden(tenant_id: str, tokens: list[str]) -> list[str]: + raise AssertionError(f"{tenant_id} has no secret to decrypt, got {tokens}") + + monkeypatch.setattr(provider_configs, "batch_decrypt_token", forbidden) + + config = decode_config(WORKSPACE, PhoenixProvider, {"api_key": None, "project": "dify"}) + + assert isinstance(config, PhoenixConfig) + assert config.api_key is None + + +@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): + ProviderConfigGateway(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 = ProviderConfigGateway(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 = ProviderConfigGateway(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_recovery.py b/api/tests/unit_tests/services/tracing/test_recovery.py index 775662e814e13b..d3b3abf6415a80 100644 --- a/api/tests/unit_tests/services/tracing/test_recovery.py +++ b/api/tests/unit_tests/services/tracing/test_recovery.py @@ -1,37 +1,96 @@ -"""Recover lost terminal notifications using real business rows and TraceTask.""" +"""Recover lost terminal notifications from real business rows through the trace assembler.""" import json -from collections.abc import Callable +from collections.abc import Callable, Generator, Iterator from dataclasses import dataclass from datetime import datetime, timedelta from decimal import Decimal -from typing import Literal +from typing import Literal, cast, override import pytest -from sqlalchemy import Engine +from cachetools import TTLCache +from dify_agent.protocol import RunEvent +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.errors import EgressError +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent +from dify_trace_sdk.ids import span_id_for +from dify_trace_sdk.model import Trace +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 events.event_handlers import update_provider_when_message_created as receiver from extensions.ext_database import db +from extensions.tracing import 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.enums import ConversationFromSource, CreatorUserRole, MessageStatus, WorkflowRunTriggeredFrom -from models.model import App, AppMode, Message +from models.account import Account, TenantAccountJoin +from models.application_trace import ApplicationTraceInvocation, ApplicationTraceRun +from models.enums import CreatorUserRole, MessageStatus, WorkflowRunTriggeredFrom +from models.model import App, AppMode, Conversation, Message 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.application_trace_repository import TraceCollectionError +from services.tracing import recovery, runtime_bridge +from services.tracing.assembly import SourceNotReadyError +from services.tracing.emitter import TraceScopes from services.tracing.service import ApplicationTracingService -from tests.helpers.unit.model_factories import make_account, make_conversation +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.apps.agent_app.test_app_runner import ( + _CancelAndWaitFailingClient, + _ExplodingSessionStore, + _FakeQueueManager, + _FakeSessionStore, + _runner, + _soul, + _StoppedQueueManager, + _UsageCancellationClient, +) +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_accepted as accepted 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_install_delivery as install_delivery, +) 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_tracing_service as tracing_service, +) from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection +from tests.unit_tests.services.tracing.test_service import ( + add_app_owner, + attach_conversation, + message_record, + sent_traces, +) @dataclass @@ -40,21 +99,51 @@ class RecoveryCase: service: ApplicationTracingService now: datetime scheduled: list[str] + destination: RecordingEgress -def attach_conversation(case: Collection) -> None: - 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=AppMode.AGENT, - inputs={}, - from_source=ConversationFromSource.CONSOLE, - ) - ) +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") + + +def recording_service( + collection: Collection, destination: RecordingEgress, scheduled: list[str], monkeypatch: pytest.MonkeyPatch +) -> ApplicationTracingService: + """The ledger the runtime hooks reach; scheduled exports are recorded instead of enqueued.""" + service = tracing_service(collection, destination) + monkeypatch.setattr(service, "schedule", scheduled.append) + monkeypatch.setattr(runtime_bridge, "_service", lambda: service) + return service + + +def tracing_application(collection: Collection) -> Flask: + """An application bound like the composition root, with events kept in memory.""" + application = Flask(__name__) + bind_tracing( + application, + scopes=TraceScopes(emitter=_Recorder(), session_factory=collection.sessions, enterprise_enabled=lambda: False), + hooks=runtime_bridge.RuntimeApplicationTracingHooks(), + ) + return application @pytest.fixture @@ -62,13 +151,12 @@ 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) + attach_conversation(collection, AppMode.AGENT) invocation = collection.invocation() ingest(collection, batch_for(collection.run, invocation)) 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 +169,12 @@ 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) + install_delivery(monkeypatch) + destination = accepted("provider-id") scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) - return RecoveryCase(collection, service, now, scheduled) + service = recording_service(collection, destination, scheduled, monkeypatch) + return RecoveryCase(collection, service, now, scheduled, destination) def complete_metadata() -> str: @@ -110,35 +197,9 @@ 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, - ) - +) -> list[TraceEvent]: + """Run the real consumer save against the owned SQL rows; returns the trace events it emitted.""" 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) def send(message: Message, *, application_generate_entity: ChatAppGenerateEntity) -> None: if notify is not None: @@ -157,8 +218,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 +238,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( @@ -251,6 +320,7 @@ 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"] @@ -270,45 +340,31 @@ def test_stopped_consumer_cannot_freeze_a_terminal_before_the_runtime_worker( outcome: Literal["succeeded", "failed", "cancelled"], delivery: str, ) -> 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 == "accepted" + else EgressError("acknowledgement lost", request_sent=True) + ], + grpc_target=None, + ) + service = recording_service(case.collection, destination, case.scheduled, monkeypatch) # 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(service, case.collection) queued = save_agent_stop(case, monkeypatch, answer) - assert case.service.capture_terminal(info) is True + assert service.capture_terminal(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 recovery.recover_terminal(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 == [] + 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, @@ -317,35 +373,33 @@ def export(_provider: str, _config: dict[str, object], trace: CanonicalTrace) -> ) 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") + trace = service.assemble(terminal) + assert trace.root.attributes["message_execution_state"] == outcome + assert trace.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) - assert len(exported) == 1 - assert exported[0] == trace + service.deliver(stored.id) + service.deliver(stored.id) + assert sent_traces(destination) == [trace] frozen = case.collection.repository.get(stored.id) assert frozen is not None assert frozen.state == delivery - 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="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(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 +413,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 +426,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) @@ -397,9 +449,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 @@ -434,26 +483,6 @@ 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) pending_invocation = None @@ -498,9 +527,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.collection).app_context(), pytest.raises(GenerateTaskStoppedError): runner._consume_stream( client.run_id, dify_context=context, @@ -527,31 +554,23 @@ def unexpected_cancel(*_args: object, **_kwargs: object) -> None: 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 - assert root.status.value == ("ok" if runtime_outcome == "succeeded" else "error") + assert trace.root.attributes["message_execution_state"] == runtime_outcome + assert trace.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 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 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 == "accepted" + assert accepted_run.attempts == 1 with case.collection.sessions() as session: message = session.get(Message, run.source_id) assert message is not None @@ -563,8 +582,6 @@ def export(_provider: str, _config: dict[str, object], payload: CanonicalTrace) 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) save_agent_stop(case, monkeypatch, "partial answer") if incomplete: @@ -586,26 +603,6 @@ def test_missing_runtime_stop_terminal_converges_without_inventing_success( 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 +644,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.collection).app_context(), pytest.raises(GenerateTaskStoppedError): runner.run( dify_context=context, agent_id=identifier("stop-agent"), @@ -672,9 +667,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 +677,7 @@ 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 +712,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.collection).app_context(): with pytest.raises(GenerateTaskStoppedError if stop_during_cleanup else RuntimeError): runner._consume_stream( client.run_id, @@ -766,9 +737,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,33 +747,30 @@ 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 == "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) + attach_conversation(collection, AppMode.AGENT) + scheduled: list[str] = [] + destination = RecordingEgress([], grpc_target=None) + stopped = RecoveryCase( + collection, + recording_service(collection, destination, scheduled, monkeypatch), + case.now, + scheduled, + destination, + ) apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_COLLECTION_TIMEOUT=60) save_agent_stop(stopped, monkeypatch, "") monkeypatch.setattr(recovery, "naive_utc_now", lambda: case.now + timedelta(seconds=120)) @@ -836,9 +803,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,8 +817,6 @@ 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"): case.service.capture_agent_app_terminal( tenant_id=identifier("foreign-tenant"), @@ -866,7 +830,7 @@ 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_assembler(case: RecoveryCase) -> None: with case.collection.sessions.begin() as session: message = session.get(Message, case.collection.run.source_id) assert message is not None @@ -888,29 +852,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.provider_trace_id == "provider-id" assert stored.attempts == 1 @@ -977,24 +933,27 @@ def test_explicit_message_failure_can_recover_without_fabricating_usage(case: Re 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( +def test_cross_tenant_source_is_rejected_before_the_assembler_runs( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch ) -> 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] = [] + called: list[TraceEvent] = [] - def forbidden(_task: TraceTask) -> None: - called.append(True) + def forbidden(event: TraceEvent) -> None: + called.append(event) - monkeypatch.setattr(TraceTask, "execute", forbidden) + monkeypatch.setattr(case.service.assembler, "assemble", forbidden) assert recovery.recover_terminal(case.service, case.collection.run) is False assert called == [] stored = case.collection.repository.get(case.collection.run.id) @@ -1002,6 +961,30 @@ def forbidden(_task: TraceTask) -> None: assert stored.error_code == "tracing_source_not_found" +def test_source_not_ready_for_assembly_is_retried_later(case: RecoveryCase, monkeypatch: pytest.MonkeyPatch) -> 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) + events: list[TraceEvent] = [] + + def not_ready(event: TraceEvent) -> None: + events.append(event) + raise SourceNotReadyError + + monkeypatch.setattr(case.service.assembler, "assemble", not_ready) + assert recovery.recover_terminal(case.service, case.collection.run) is False + assert [event.payload for event in events] == [ + MessageCompleted(message_id=case.collection.run.source_id, conversation_id=identifier("conversation-recover")) + ] + assert events[0].correlation.session_id == "session-original" + stored = case.collection.repository.get(case.collection.run.id) + assert stored is not None + assert (stored.state, stored.error_code, stored.outer_trace) == ("collecting", None, None) + assert case.scheduled == [] + + def test_unresolved_recovery_rotates_behind_unchecked_sources(case: RecoveryCase) -> None: other = seed(case.collection.sessions, "other-recovery") assert case.collection.repository.pending_ids(limit=1) == [case.collection.run.id] @@ -1023,6 +1006,48 @@ def test_no_final_marker_or_paused_message_is_not_recovered(case: RecoveryCase, 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) -> ApplicationTraceRun: + 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 + return run + + @pytest.mark.parametrize( ("status", "finished"), [ @@ -1038,44 +1063,9 @@ 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, @@ -1087,10 +1077,40 @@ def test_workflow_recovery_requires_committed_terminal_state( 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) == ("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 recovery.recover_terminal(case.service, run) 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 +1118,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 +1138,14 @@ 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) 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 +1181,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 +1197,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 +1222,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 = "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) 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..71f36e1b1e2ebc --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_registry.py @@ -0,0 +1,56 @@ +from collections.abc import Iterator +from importlib.metadata import EntryPoint + +import pytest + +from services.tracing import registry +from services.tracing.registry import ( + ENTRY_POINT_GROUP, + SUPPORTED_PROVIDERS, + TraceProviderUnavailableError, + UnknownTraceProviderError, + installed_providers, + provider_class, +) + + +@pytest.fixture(autouse=True) +def fresh_registry() -> Iterator[None]: + provider_class.cache_clear() + yield + provider_class.cache_clear() + + +def test_every_supported_provider_is_installed_under_its_own_id() -> None: + assert installed_providers() == SUPPORTED_PROVIDERS + for provider_id in SUPPORTED_PROVIDERS: + assert provider_class(provider_id).id == provider_id + + +def test_unknown_providers_are_client_errors() -> None: + with pytest.raises(UnknownTraceProviderError): + provider_class("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): + provider_class("langfuse") + assert "langfuse" not in installed_providers() 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 ebff6db3caca43..4b2409bebe6305 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..ed931da6c63a42 100644 --- a/api/tests/unit_tests/services/tracing/test_runtime_bridge.py +++ b/api/tests/unit_tests/services/tracing/test_runtime_bridge.py @@ -2,15 +2,22 @@ from copy import deepcopy from dataclasses import dataclass, field, replace from datetime import UTC, datetime, timedelta -from typing import cast +from typing import Literal, cast from uuid import UUID, uuid4 -import httpx import pytest from dify_agent.layers.execution_context import ApplicationTracingContext from dify_agent.protocol.tracing import AgentTraceBatch, AgentTraceSpan -from flask import Flask -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest +from dify_trace_sdk.events import ( + ConversationNamed, + MessageCompleted, + ModerationChecked, + TraceCorrelation, + TraceEvent, + TracePayload, + WorkflowRunCompleted, +) +from dify_trace_sdk.testing import RecordingEgress, response from sqlalchemy import delete, func, select from sqlalchemy.exc import OperationalError from sqlalchemy.orm import Session, sessionmaker @@ -18,20 +25,8 @@ 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 graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus from graphon.runtime import VariablePool from models.agent import ( @@ -44,36 +39,48 @@ ) from models.agent_config_entities import AgentSoulConfig, WorkflowNodeJobConfig from models.application_trace import ApplicationTraceRun -from models.enums import ConversationFromSource, CreatorUserRole +from models.enums import CreatorUserRole from models.enums import WorkflowRunTriggeredFrom as _application_trace_handoff_WorkflowRunTriggeredFrom from models.model import App, AppMode, Message, TraceAppConfig from models.workflow import ( + WorkflowAppLog, + WorkflowAppLogCreatedFrom, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun, ) -from models.workflow import ( - WorkflowRunTriggeredFrom as _workflow_tool_tracing_boundary_WorkflowRunTriggeredFrom, -) +from models.workflow import WorkflowRunTriggeredFrom as _workflow_tool_tracing_boundary_WorkflowRunTriggeredFrom from models.workflow import WorkflowType as _application_trace_handoff_WorkflowType -from repositories.application_trace_repository import ApplicationTraceRepository +from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError from services.tracing import runtime_bridge +from services.tracing.assembly import TraceAssembler +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.delivery import TraceCollectorUnavailableError, TraceDeliveryService from services.tracing.normalization import remap_outer_trace -from services.tracing.provider_dispatch import export_trace +from services.tracing.records import WorkflowTraceInfo 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.helpers.unit.model_factories import make_account, make_app, 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_no_egress as no_egress from tests.unit_tests.services.tracing.test_service import ( _application_trace_collection_seed as _application_trace_handoff_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 ( + _application_trace_collection_tracing_service as tracing_service, +) from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection +from tests.unit_tests.services.tracing.test_service import ( + add_app_owner, + attach_conversation, + install_decryption, + trace_event, +) -"""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,6 +89,9 @@ } +STARTED_AT = datetime(2026, 10, 15, tzinfo=UTC) + + @dataclass class Case: collection: Collection @@ -93,69 +103,79 @@ def setup( sessions: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, name: str, mode: AppMode, *, snapshot: bool ) -> Case: apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + install_decryption(monkeypatch) collection = _application_trace_handoff_seed(sessions, name) + attach_conversation(collection, mode) with sessions.begin() as session: app = session.get(App, collection.run.app_id) - 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, - mode=mode, - inputs={}, - from_source=ConversationFromSource.CONSOLE, - ) - ) if not snapshot: session.execute(delete(ApplicationTraceRun).where(ApplicationTraceRun.id == collection.run.id)) - service = ApplicationTracingService(collection.repository) + service = tracing_service(collection, no_egress()) scheduled: list[str] = [] monkeypatch.setattr(service, "schedule", scheduled.append) monkeypatch.setattr(runtime_bridge, "_service", lambda: service) return Case(collection, service, scheduled) -def message_info(case: Case) -> MessageTraceInfo: +def deliver(case: Case, event: TraceEvent, destination: RecordingEgress) -> None: + """Deliver one event the way the trace worker does, with this case's ledger as the Agent collector.""" + TraceDeliveryService( + session_factory=case.collection.sessions, + assembler=case.service.assembler, + builder=CanonicalTraceBuilder(), + egress=destination, + enterprise_sink=None, + collector=case.service, + max_export_attempts=1, + ).deliver(event, enterprise_delivered=False) + + +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 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, 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( @@ -209,7 +229,7 @@ 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: metadata: dict[str, object] = {"app_id": row.app_id} if parent is not None: metadata["parent_trace_context"] = { @@ -218,6 +238,10 @@ def workflow_info(row: WorkflowRun, parent: WorkflowRun | None = None) -> Workfl } 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 +250,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, @@ -235,9 +263,14 @@ def workflow_info(row: WorkflowRun, parent: WorkflowRun | None = None) -> Workfl ) +def run_count(case: Case) -> int | None: + with case.collection.sessions() as session: + return session.scalar(select(func.count()).select_from(ApplicationTraceRun)) + + class TestApplicationTraceHandoffCases: @staticmethod - def test_chatflow_message_event_cannot_export_again_via_legacy( + def test_chatflow_message_event_is_not_exported_again( sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch ) -> None: case = setup(sqlite_session_factory, monkeypatch, "chatflow", AppMode.ADVANCED_CHAT, snapshot=False) @@ -247,16 +280,17 @@ def test_chatflow_message_event_cannot_export_again_via_legacy( 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) == 1 @staticmethod @pytest.mark.parametrize("mode", [AppMode.AGENT, AppMode.ADVANCED_CHAT]) @@ -287,7 +321,9 @@ def fail_start(**_kwargs: object) -> None: backend_binding_ref="backend", ) assert runtime_bridge.prepare_agent_app(context) is None - assert dispatch(case, message_info(case), monkeypatch) == [] + destination = no_egress() + deliver(case, message_event(case), destination) + assert destination.requests == [] stored = case.collection.repository.source_run( tenant_id=case.collection.run.tenant_id, app_id=case.collection.run.app_id, @@ -306,83 +342,140 @@ def test_auxiliary_events_resolve_owner_without_bypassing_content_policy( sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, 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", + 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, payload), destination) + assert len(destination.requests) == (1 if mode is AppMode.CHAT else 0) assert case.scheduled == [] + assert run_count(case) == 0 @staticmethod - def test_plain_workflow_preserves_legacy_parent_context( + def test_plain_nested_workflow_keeps_its_parent_context_for_the_provider( sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch ) -> None: parent_case = setup(sqlite_session_factory, monkeypatch, "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, _ = 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 = child_case.service.assembler.assemble(event) + assert isinstance(record, WorkflowTraceInfo) + assert child_case.service.capture_terminal(record) is False + assert child_case.service.link_managed_parent(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(child_case) == 0 @staticmethod - def test_plain_chatflow_auxiliary_events_keep_the_legacy_route( + def test_plain_chatflow_auxiliary_events_keep_the_provider_route( sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch ) -> None: case = setup(sqlite_session_factory, monkeypatch, "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, payload), destination) + assert len(destination.requests) == 1 @staticmethod - def test_legacy_child_of_managed_parent_uses_explicit_link_not_unavailable_context( + def test_provider_exported_child_of_managed_parent_names_the_collected_parent( sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch ) -> None: parent_case = setup(sqlite_session_factory, monkeypatch, "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, monkeypatch, "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 + linked = child_case.service.link_managed_parent(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], monkeypatch: pytest.MonkeyPatch + ) -> None: + case = setup(sqlite_session_factory, monkeypatch, "collector-unavailable", AppMode.AGENT, snapshot=True) + + def unavailable(**_kwargs: object) -> None: + raise OperationalError("statement", {}, Exception("sensitive database error")) + + monkeypatch.setattr(case.collection.repository, "source_run", unavailable) + destination = no_egress() + with pytest.raises(TraceCollectorUnavailableError): + deliver(case, message_event(case), destination) + assert destination.requests == [] @staticmethod def test_mixed_workflow_exports_normal_llm_and_agent_calls_without_double_usage( @@ -391,13 +484,9 @@ def test_mixed_workflow_exports_normal_llm_and_agent_calls_without_double_usage( case = setup(sqlite_session_factory, monkeypatch, "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] = ( @@ -484,29 +573,23 @@ def test_mixed_workflow_exports_normal_llm_and_agent_calls_without_double_usage( 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 = case.service.assembler.assemble(workflow_event(row, None)) + assert isinstance(record, WorkflowTraceInfo) + assert case.service.capture_terminal(record) is True + assert case.scheduled == [managed.id] + destination = exported() + tracing_service(Collection(case.collection.repository, case.collection.sessions, managed), destination).deliver( + managed.id + ) + delivered = case.collection.repository.get(managed.id) + assert delivered is not None + assert delivered.state == "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.""" @@ -602,7 +685,12 @@ class TestRuntimeBridgeCases: @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)) + target = ApplicationTracingService( + ApplicationTraceRepository(sqlite_session_factory), + assembler=TraceAssembler(sqlite_session_factory, files_url="https://files.test", enterprise_details=False), + builder=CanonicalTraceBuilder(), + egress=no_egress(), + ) case = BridgeCase(target) def prepare(**values: object) -> ApplicationTracingContext: @@ -616,11 +704,6 @@ def start(**values: object) -> 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 @@ -628,7 +711,6 @@ def agent_terminal(**values: object) -> None: 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 @@ -652,8 +734,9 @@ def forbidden() -> ApplicationTracingService: conversation_id=None, graph_data={}, ) - assert runtime_bridge.capture_terminal(cast(BaseTraceInfo, object())) is False - assert runtime_bridge.queue_enabled_at_start("app") is False + runtime_bridge.capture_agent_app_terminal( + tenant_id="tenant", app_id="app", message_id="message", outcome="succeeded", output="answer" + ) assert calls == [] @staticmethod @@ -722,12 +805,6 @@ def test_business_hook_failure_logs_type_without_sensitive_exception_message( 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( @@ -758,8 +835,6 @@ def test_graph_string_with_agent_is_recognized() -> None: @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( tenant_id="tenant", app_id="agent-app", @@ -784,8 +859,6 @@ def test_worker_terminal_handoff_preserves_outcome_without_sse(service: BridgeCa def test_terminal_handoff_gate_and_fail_open_log_redact_sensitive_details( service: BridgeCase, monkeypatch: pytest.MonkeyPatch, 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 @@ -818,35 +891,18 @@ 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() -> None: + hooks = runtime_bridge.RuntimeApplicationTracingHooks() + assert hooks.prepare_agent_app is runtime_bridge.prepare_agent_app + assert hooks.prepare_workflow_agent is runtime_bridge.prepare_workflow_agent + assert hooks.freeze_workflow_run is runtime_bridge.freeze_workflow_run + assert hooks.capture_agent_app_terminal is runtime_bridge.capture_agent_app_terminal + assert {name for name in dir(hooks) if not name.startswith("_")} == { + "prepare_agent_app", + "prepare_workflow_agent", + "freeze_workflow_run", + "capture_agent_app_terminal", + } """Hidden Workflow Tool Agents must not invent a trace owned by the source app.""" @@ -968,7 +1024,12 @@ def _workflow_tool_tracing_boundary_seed( ) node_record.created_at = now session.add(node_record) - service = ApplicationTracingService(ApplicationTraceRepository(sessions)) + service = ApplicationTracingService( + ApplicationTraceRepository(sessions), + assembler=TraceAssembler(sessions, files_url="https://files.test", enterprise_details=False), + builder=CanonicalTraceBuilder(), + egress=no_egress(), + ) apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) monkeypatch.setattr(runtime_bridge, "_service", lambda: service) entity = _workflow_tool_tracing_boundary_workflow_entity() @@ -1086,6 +1147,10 @@ def test_hidden_child_skip_keeps_normal_root_agent_trace_assemblable( ) 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 +1159,28 @@ 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) + outer = remap_outer_trace(outer, trace_id=run.trace_id, capture_content=False, synthetic_span_ids=frozenset()) + 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) 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) + span_ids = {span.span_id for span in trace.spans} + assert all(span.parent_span_id is None or span.parent_span_id in span_ids for span in trace.spans) 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_service.py b/api/tests/unit_tests/services/tracing/test_service.py index 9bcc701f97be8b..0354bff549d33a 100644 --- a/api/tests/unit_tests/services/tracing/test_service.py +++ b/api/tests/unit_tests/services/tracing/test_service.py @@ -1,29 +1,60 @@ import json +import threading +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass from datetime import UTC, datetime, timedelta from decimal import Decimal +from typing import 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 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 machinery.context import RequestContext +from models.account import Tenant, TenantAccountJoin, TenantAccountRole from models.application_trace import ApplicationTraceBatch, ApplicationTraceInvocation, ApplicationTraceRun from models.enums import ConversationFromSource, CreatorUserRole from models.model import App, AppMode, Message, TraceAppConfig -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom +from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun +from repositories import application_trace_repository as repository_module +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError +from services.app_tracing_config_service import AppTracingConfigService +from services.tracing import provider_configs as provider_configs_module from services.tracing import service as service_module -from services.tracing.provider_dispatch import TraceExportResult, TraceExportUncertainError +from services.tracing.assembly import TraceAssembler +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.provider_configs import ProviderConfigGateway +from services.tracing.records import GenerateNameTraceInfo, MessageTraceInfo +from services.tracing.registry import UnknownTraceProviderError, provider_class 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, +) """Exercise collection transactions and delivery state against real SQLite rows.""" @@ -48,6 +79,88 @@ def invocation(self, name: str = "first") -> ApplicationTraceInvocation: ) +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 install_decryption(monkeypatch: pytest.MonkeyPatch) -> list[str]: + """A stored ``encrypted-x`` secret decrypts to ``plaintext-x``; returns the ciphertexts decrypted.""" + ciphertexts: list[str] = [] + + def decrypt(_tenant_id: str, tokens: list[str]) -> list[str]: + ciphertexts.extend(tokens) + return [token.replace("encrypted-", "plaintext-", 1) for token in tokens] + + monkeypatch.setattr(provider_configs_module, "batch_decrypt_token", decrypt) + return ciphertexts + + +def _application_trace_collection_install_delivery(monkeypatch: pytest.MonkeyPatch) -> list[str]: + """Deliver Agent traces through ``ScriptedProvider``; returns the ciphertexts decrypted.""" + + def provider_class(provider_id: str) -> type[ScriptedProvider]: + if provider_id != ScriptedProvider.id: + raise UnknownTraceProviderError(provider_id) + return ScriptedProvider + + monkeypatch.setattr(service_module, "provider_class", provider_class) + return install_decryption(monkeypatch) + + +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_tracing_service( + case: _ApplicationTraceCollectionCollection, egress: Egress +) -> ApplicationTracingService: + return ApplicationTracingService( + case.repository, + assembler=TraceAssembler(case.sessions, files_url="https://files.test", enterprise_details=False), + builder=CanonicalTraceBuilder(), + egress=egress, + ) + + def _application_trace_collection_seed( sessions: sessionmaker[Session], name: str, @@ -97,9 +210,10 @@ 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", }, ) ) @@ -111,6 +225,63 @@ def _application_trace_collection_seed( 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(service: ApplicationTracingService, 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 = service.assembler.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, @@ -147,17 +318,24 @@ def _application_trace_collection_ingest( 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) + collection.repository.terminal(run.id, trace.model_dump(mode="json"), root.span_id) def _application_trace_collection_claim( @@ -184,7 +362,10 @@ def test_batch_duplicate_is_idempotent_after_final_and_acceptance( 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") + collection.run.id, + leased[1], + state="accepted", + result=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) @@ -327,7 +508,12 @@ def test_expired_delivery_lease_accepts_confirmed_late_ack_for_same_attempt( 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")) + collection.repository.complete( + collection.run.id, + leased[1], + state="accepted", + result=SendResult(trace_id="late", trace_url=None, span_urls={}), + ) stored = collection.repository.get(collection.run.id) assert stored is not None assert stored.state == "accepted" @@ -377,42 +563,39 @@ 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"), + [ + ("accepted", response(200, {"trace_id": "vendor-trace", "trace_url": "https://trace.test/example"})), + ("uncertain", EgressError("acknowledgement lost", request_sent=True)), + ("retrying", response(503, {})), + ], + ) def test_delivery_checkpoint_controls_replays( - collection: _ApplicationTraceCollectionCollection, monkeypatch: pytest.MonkeyPatch, outcome: str + collection: _ApplicationTraceCollectionCollection, + monkeypatch: pytest.MonkeyPatch, + outcome: str, + 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) + ciphertexts = _application_trace_collection_install_delivery(monkeypatch) + egress = RecordingEgress([reply], grpc_target=None) + service = _application_trace_collection_tracing_service(collection, egress) 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 ciphertexts == ["encrypted-public", "encrypted-original"] stored = collection.repository.get(collection.run.id) assert stored is not None assert stored.state == 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: @@ -445,14 +628,13 @@ def test_disabled_source_decision_stays_disabled_after_config_enable( ) -> 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) + attach_conversation(case, AppMode.AGENT) + service = _application_trace_collection_tracing_service(case, _application_trace_collection_no_egress()) assert case.run.state == "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, @@ -463,17 +645,7 @@ def test_disabled_source_decision_stays_disabled_after_config_enable( ) 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(message_record(service, case)) is True stored = case.repository.get(case.run.id) assert stored is not None assert stored.state == "not_collected" @@ -486,7 +658,8 @@ def test_capture_agent_disabled_preserves_outer_only_content_policy( ) -> 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) + attach_conversation(case, AppMode.AGENT) + service = _application_trace_collection_tracing_service(case, _application_trace_collection_no_egress()) scheduled: list[str] = [] monkeypatch.setattr(service, "schedule", scheduled.append) assert ( @@ -503,27 +676,14 @@ def test_capture_agent_disabled_preserves_outer_only_content_policy( ) 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", - ) + info = message_record(service, case).model_copy(update={"inputs": "private input", "outputs": "private output"}) assert service.capture_terminal(info) is True assert 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 == [] @@ -532,21 +692,16 @@ def test_capture_agent_disabled_preserves_outer_only_content_policy( 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(_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 @@ -565,8 +720,19 @@ def test_deleting_selected_configuration_keeps_historical_destination( _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_install_delivery = _application_trace_collection_install_delivery +_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 +_application_trace_service_boundaries_tracing_service = _application_trace_collection_tracing_service + + +FROZEN_LANGFUSE: dict[str, object] = { + "public_key": "encrypted-public", + "secret_key": "encrypted-original", + "host": "https://langfuse.test", + "project_key": "original", +} def source(case: _ApplicationTraceServiceBoundariesCollection) -> dict[str, str]: @@ -578,6 +744,21 @@ def source(case: _ApplicationTraceServiceBoundariesCollection) -> dict[str, str] } +def forbid_credentials(monkeypatch: pytest.MonkeyPatch, reason: str) -> None: + def forbidden(*_args: object, **_kwargs: object) -> None: + pytest.fail(reason) + + monkeypatch.setattr(provider_configs_module, "batch_decrypt_token", forbidden) + + +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( @@ -585,7 +766,9 @@ def test_prepare_invocation_returns_only_opaque_context_and_respects_capture_gat ) -> 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 = _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ) context = service.prepare_invocation( tenant_id=case.run.tenant_id, app_id=case.run.app_id, @@ -625,7 +808,9 @@ def test_status_denies_foreign_or_absent_source_without_destination_details( _application_trace_service_boundaries_identifier("foreign") ) with pytest.raises(TraceCollectionError) as caught: - ApplicationTracingService(case.repository).status(**identity) + _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ).status(**identity) assert caught.value.code == "tracing_source_not_found" assert caught.value.status == 404 assert "encrypted-original" not in str(caught.value) @@ -639,13 +824,8 @@ def test_status_uses_frozen_destination_and_selects_latest_node_attempt( 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"} + forbid_credentials(monkeypatch, "status must not decrypt credentials") node_id = _application_trace_service_boundaries_identifier("selected-node") invocation = case.repository.invocation( run_id=case.run.id, @@ -664,21 +844,27 @@ def forbidden(*_args: object, **_kwargs: object) -> None: case.run.id, lease[1], state="accepted", - result=TraceExportResult( - "provider-trace", + result=SendResult( + trace_id="provider-trace", trace_url="https://frozen.example/trace", span_urls={invocation.span_id: "https://frozen.example/span"}, ), ) - service = ApplicationTracingService(case.repository) + service = _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ) result = service.status(**source(case), node_execution_id=node_id) assert result["status"] == "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["queryable"] is None assert result["capture_content"] is False + capabilities = result["capabilities"] + assert isinstance(capabilities, dict) + assert capabilities["provider"] == "langfuse" assert "provider_config" not in result assert "encrypted-original" not in str(result) with pytest.raises(TraceCollectionError, match="tracing_node_not_found"): @@ -695,7 +881,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( @@ -721,7 +907,9 @@ def test_new_run_freezes_updated_project_without_changing_previous_run( session_id="new-session", ) assert new_run is not None - service = ApplicationTracingService(case.repository) + service = _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ) assert service.status(**source(case))["project_url"] == "https://langfuse.test/project/original" assert ( service.status(**{**source(case), "source_id": new_id})["project_url"] == "https://new.example/project/new" @@ -729,53 +917,63 @@ def test_new_run_freezes_updated_project_without_changing_previous_run( @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)) + freeze_config(case, "langfuse", frozen) + result = _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ).status(**source(case)) 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"}) + destination = provider_class(provider) + expected = destination(_application_trace_service_boundaries_no_egress()).project_url( + destination.config_model.model_validate(frozen) + ) + result = _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ).status(**source(case)) + assert expected is not None + assert result["project_url"] == expected + capabilities = result["capabilities"] + assert isinstance(capabilities, dict) + assert capabilities["provider"] == provider @staticmethod @pytest.mark.parametrize("disabled", [True, False]) def test_uncollected_status_and_missing_run_retry_do_not_construct_provider( - sqlite_session_factory: sessionmaker[Session], disabled: bool + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, disabled: bool ) -> None: case = _application_trace_service_boundaries_seed(sqlite_session_factory, "missing", enabled=not disabled) if not disabled: @@ -783,7 +981,14 @@ 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) + + def forbidden(_provider_id: str) -> None: + pytest.fail("an uncollected source has no destination to construct") + + monkeypatch.setattr(service_module, "provider_class", forbidden) + service = _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ) assert service.status(**source(case)) == { "status": "not_collected", "source_kind": "message", @@ -806,7 +1011,9 @@ def test_retry_acknowledgement_guards_schedule_and_keeps_frozen_snapshot( 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) + service = _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ) monkeypatch.setattr(service, "schedule", scheduled.append) with pytest.raises(TraceCollectionError, match="tracing_duplicate_acknowledgement_required"): service.retry(**source(case), acknowledge_possible_duplicates=False) @@ -821,40 +1028,93 @@ def test_retry_acknowledgement_guards_schedule_and_keeps_frozen_snapshot( @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, "failed", "provider_retry_exhausted"), + (ExportRetryableError("private-transport-error"), 2, "retrying", "provider_temporarily_unavailable"), + (ExportRejectedError("private-rejection"), 2, "failed", "provider_export_rejected"), + (ExportOutcomeUnknownError("private-acknowledgement"), 2, "uncertain", "provider_acceptance_unknown"), + (ValueError("private-credential-error"), 2, "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: str, code: str, ) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_MAX_ATTEMPTS=1) + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_MAX_ATTEMPTS=max_attempts) 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) + _application_trace_service_boundaries_install_delivery(monkeypatch) + destination = RecordingEgress([failure], grpc_target=None) + service = _application_trace_service_boundaries_tracing_service(case, destination) 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.next_attempt_at is not None) is (state == "retrying") + assert current.attempts == 1 + 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], + monkeypatch: pytest.MonkeyPatch, + 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) + _application_trace_service_boundaries_install_delivery(monkeypatch) + if unreadable == "decryption": + + def undecryptable(_tenant_id: str, _tokens: list[str]) -> list[str]: + raise ValueError("private-ciphertext") + + monkeypatch.setattr(provider_configs_module, "batch_decrypt_token", undecryptable) + else: + freeze_config(case, "langfuse", {**FROZEN_LANGFUSE, "host": "private-host-without-scheme"}) + destination = RecordingEgress([], grpc_target=None) + _application_trace_service_boundaries_tracing_service(case, destination).deliver(case.run.id) + current = case.repository.get(case.run.id) + assert current is not None + assert (current.state, current.error_code) == ("failed", "tracing_config_unreadable") + assert destination.requests == [] + assert "private-" not in caplog.text + with pytest.raises(TraceCollectionError, match="tracing_retry_not_allowed"): + case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) + + @staticmethod + def test_unavailable_frozen_provider_fails_before_decryption( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> 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) + forbid_credentials(monkeypatch, "an unavailable provider must not decrypt credentials") + _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ).deliver(case.run.id) + current = case.repository.get(case.run.id) + assert current is not None + assert (current.state, current.error_code) == ("failed", "tracing_provider_unavailable") + @staticmethod @pytest.mark.parametrize("capture_content", [False, True]) @pytest.mark.parametrize( @@ -906,45 +1166,25 @@ 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) + _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ).settle_preview(run) current = case.repository.get(run.id) assert current is not None if status == WorkflowNodeExecutionStatus.RUNNING: assert current.outer_trace is None else: assert current.outer_trace is not None - trace = service_module.CanonicalTrace.model_validate(current.outer_trace) - root = trace.spans[0] + 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,6 +1192,7 @@ 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"), ], ) @@ -965,11 +1206,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(), @@ -982,10 +1227,11 @@ def test_invalid_runtime_tree_is_rejected_before_provider_credentials_or_io( 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) + monkeypatch.setattr(provider_configs_module, "batch_decrypt_token", forbidden) + monkeypatch.setattr(service_module, "provider_class", forbidden) + _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ).deliver(case.run.id) current = case.repository.get(case.run.id) assert current is not None assert current.state == "failed" @@ -995,7 +1241,40 @@ def forbidden(*_args: object, **_kwargs: object) -> None: def test_assemble_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) + _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ).assemble(case.run) + + @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) + stored = case.repository.get(case.run.id) + assert stored is not None + trace = _application_trace_service_boundaries_tracing_service( + case, _application_trace_service_boundaries_no_egress() + ).assemble(stored) + 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,8 +1282,10 @@ 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 +_invalid_trace_delivery_tracing_service = _application_trace_collection_tracing_service class TestInvalidTraceDeliveryCases: @@ -1039,9 +1320,9 @@ def test_invalid_tree_is_permanent_before_decryption_or_provider_io( 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) + monkeypatch.setattr(provider_configs_module, "batch_decrypt_token", forbidden) + monkeypatch.setattr(service_module, "provider_class", forbidden) + _invalid_trace_delivery_tracing_service(case, _invalid_trace_delivery_no_egress()).deliver(case.run.id) current = case.repository.get(case.run.id) assert current is not None assert current.state == "failed" @@ -1049,3 +1330,669 @@ def forbidden(*_args: object, **_kwargs: object) -> None: 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) + + +"""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 +_application_trace_node_alias_tracing_service = _application_trace_collection_tracing_service + + +@dataclass +class NodeAlias: + case: _ApplicationTraceNodeAliasCollection + node: dict[str, object] + workflow: dict[str, object] + requests: list[str] + logstore: bool + + def service(self) -> ApplicationTracingService: + return _application_trace_node_alias_tracing_service(self.case, _application_trace_node_alias_no_egress()) + + 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=_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 + self.case.repository.complete( + self.case.run.id, + lease[1], + state="accepted", + result=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: + 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( + _ApplicationTraceNodeAliasCollection(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: + _application_trace_node_alias_terminal(alias.case) + lease = _application_trace_node_alias_claim(alias.case) + assert lease is not None + alias.case.repository.complete( + alias.case.run.id, + lease[1], + state="accepted", + result=SendResult(trace_id="outer-only", trace_url="https://provider.test/outer", span_urls={}), + ) + 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=_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(TraceCollectionError, match="tracing_node_not_found") as error: + alias.status(_application_trace_node_alias_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=_application_trace_node_alias_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] = _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(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. +""" + + +_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_install_delivery = _application_trace_collection_install_delivery +_export_lease_seed = _application_trace_collection_seed +_export_lease_terminal = _application_trace_collection_terminal +_export_lease_tracing_service = _application_trace_collection_tracing_service + + +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 result(trace_id: str) -> SendResult: + return SendResult(trace_id=trace_id, trace_url=None, span_urls={}) + + +def _stored(case: _ExportLeaseCollection) -> ApplicationTraceRun: + row = case.repository.get(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 == "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, monkeypatch: pytest.MonkeyPatch + ) -> None: + clock = [naive_utc_now()] + monkeypatch.setattr(repository_module, "naive_utc_now", lambda: clock[0]) + leased = _export_lease_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 _export_lease_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 _export_lease_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: _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 == "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 case.repository.renew_export_lease(case.run.id, "other-token") + case.repository.complete(case.run.id, "other-token", state="accepted", result=result("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: _ExportLeaseCollection) -> None: + token = _expire(case) + confirmed = SendResult( + trace_id="confirmed", + trace_url="https://example.com/trace", + span_urls={"a" * 16: "https://example.com/span"}, + ) + case.repository.complete(case.run.id, token, state="accepted", result=confirmed) + 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 == confirmed.trace_url + assert current.span_urls == {"a" * 16: "https://example.com/span"} + case.repository.complete(case.run.id, token, state="accepted", result=result("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: _ExportLeaseCollection, + ) -> 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=result("old")) + assert _stored(case).state == "ready" + second = _export_lease_claim(case) + assert second is not None + assert second[1] != old_token + case.repository.complete(case.run.id, old_token, state="accepted", result=result("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=result("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: _ExportLeaseCollection, 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: _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" + case.repository.complete(case.run.id, token, state="accepted", result=result("late")) + assert _stored(case).state == "uncertain" + + @staticmethod + @pytest.mark.parametrize("transient_failure", [False, True]) + @pytest.mark.parametrize( + ("outcome", "reply"), + [ + ("accepted", response(200, {"trace_id": "ack"})), + ("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: str, + reply: EgressResponse | Exception, + ) -> 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]) + _export_lease_install_delivery(monkeypatch) + renewed = threading.Event() + calls: list[int] = [] + 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") + renewal = original(run_id, token) + renewed.set() + return renewal + + monkeypatch.setattr(case.repository, "renew_export_lease", renew) + 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 = renewed.wait(3) + second_claim = case.repository.claim(case.run.id, collection_timeout=60, lease_seconds=1) + observed.append((heartbeat_committed, second_claim is None, _stored(case).state)) + + _export_lease_tracing_service(case, _SlowDestination([reply], while_exporting)).deliver(case.run.id) + assert observed == [(True, True, "exporting")] + 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: _ExportLeaseCollection, monkeypatch: pytest.MonkeyPatch + ) -> None: + _export_lease_install_delivery(monkeypatch) + + 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"): + _export_lease_tracing_service(case, _export_lease_accepted("ack")).deliver(case.run.id) + assert not any(thread.name == f"application-trace-lease-{case.run.id}" for thread in threading.enumerate()) + + +_application_trace_ownership_seed = _application_trace_collection_seed +_application_trace_ownership_tracing_service = _application_trace_collection_tracing_service + + +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 = _application_trace_ownership_tracing_service(case, RecordingEgress([], grpc_target=None)) + 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(record) is True + + @staticmethod + def test_runs_frozen_for_an_unavailable_provider_report_no_capabilities() -> None: + assert ApplicationTracingService._frozen_capabilities("not-a-provider") 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..b96a4e3ee5731b 100644 --- a/api/tests/unit_tests/tasks/test_application_trace_task.py +++ b/api/tests/unit_tests/tasks/test_application_trace_task.py @@ -1,50 +1,73 @@ """Exercise Celery entrypoints against durable application trace state.""" from datetime import timedelta +from uuid import UUID +import httpx import pytest +from sqlalchemy import Engine from sqlalchemy.orm import Session, sessionmaker -from core.ops.unified_trace.entities import CanonicalTrace +from core.helper import ssrf_proxy +from extensions.ext_database import db from libs.datetime_utils import naive_utc_now from models.application_trace import ApplicationTraceRun +from models.model import AppMode from services.tracing import service as service_module -from services.tracing.provider_dispatch import TraceExportResult -from services.tracing.service import ApplicationTracingService +from services.tracing.service import application_tracing_service 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_no_egress as no_egress 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 ( + _application_trace_collection_tracing_service as tracing_service, +) +from tests.unit_tests.services.tracing.test_service import attach_conversation, install_decryption, message_record -def test_worker_task_exports_once_and_replayed_task_observes_checkpoint( +@pytest.fixture +def worker_database(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> None: + """The worker's service factory binds the application engine; bind it to the test database.""" + monkeypatch.setattr(type(db), "engine", property(lambda _db: sqlite_engine)) + + +@pytest.mark.usefixtures("worker_database") +def test_worker_task_exports_once_through_the_ssrf_proxy_and_replay_observes_checkpoint( 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] = [] + install_decryption(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.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}" + + +@pytest.mark.usefixtures("worker_database") +def test_worker_records_carry_no_enterprise_names(sqlite_session_factory: sessionmaker[Session]) -> None: + case = seed(sqlite_session_factory, "worker-names") + attach_conversation(case, AppMode.AGENT) + record = message_record(application_tracing_service(), case) + assert (record.metadata["app_name"], record.metadata["workspace_name"]) == ("", "") def test_reconcile_only_schedules_due_pending_rows_and_applies_retention( @@ -64,7 +87,7 @@ def test_reconcile_only_schedules_due_pending_rows_and_applies_retention( future.next_attempt_at = naive_utc_now() + timedelta(hours=1) expired.state = "accepted" expired.created_at = naive_utc_now() - timedelta(days=4) - service = ApplicationTracingService(due.repository) + service = tracing_service(due, no_egress()) scheduled: list[str] = [] monkeypatch.setattr(service, "schedule", scheduled.append) monkeypatch.setattr(service_module, "application_tracing_service", lambda: service) 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..befcbeafea0731 --- /dev/null +++ b/api/tests/unit_tests/tasks/test_trace_delivery_task.py @@ -0,0 +1,184 @@ +from datetime import UTC, datetime, timedelta + +import pytest +from celery.exceptions import Retry +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent +from dify_trace_sdk.testing import message_trace + +from services.tracing.assembly import SourceNotFoundError, SourceNotReadyError +from services.tracing.delivery import ExportJob, TraceCollectorUnavailableError +from services.tracing.emitter import payload_key +from tasks import trace_delivery_task +from tasks.trace_delivery_task import deliver_trace_event, retry_trace_export, sweep_trace_payloads + + +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 _Delivery: + def __init__(self, failure: Exception | None) -> None: + self.failure = failure + self.delivered: list[tuple[TraceEvent, bool]] = [] + self.retried: list[ExportJob] = [] + + def deliver(self, event: TraceEvent, *, enterprise_delivered: bool) -> None: + self.delivered.append((event, enterprise_delivered)) + if self.failure is not None: + raise self.failure + + def retry(self, job: ExportJob) -> None: + self.retried.append(job) + + +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 _use_delivery(monkeypatch: pytest.MonkeyPatch, failure: Exception | None) -> _Delivery: + delivery = _Delivery(failure) + monkeypatch.setattr(trace_delivery_task, "trace_delivery_service", lambda: delivery) + return delivery + + +@pytest.mark.usefixtures("storage") +def test_inline_events_are_delivered(monkeypatch: pytest.MonkeyPatch) -> None: + delivery = _use_delivery(monkeypatch, None) + + deliver_trace_event.run({"event": _event().model_dump_json()}) + + ((event, enterprise_delivered),) = delivery.delivered + assert (event.event_id, enterprise_delivered) == ("event-1", False) + + +def test_stored_events_are_loaded_and_their_payload_removed(monkeypatch: pytest.MonkeyPatch, storage: _Storage) -> None: + delivery = _use_delivery(monkeypatch, 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.delivered) == 1 + assert key not in storage.files + + +def test_unreadable_events_are_dropped(monkeypatch: pytest.MonkeyPatch, storage: _Storage) -> None: + delivery = _use_delivery(monkeypatch, 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.delivered == [] + assert storage.files == {} + + +@pytest.mark.usefixtures("storage") +def test_missing_sources_are_dropped(monkeypatch: pytest.MonkeyPatch) -> None: + _use_delivery(monkeypatch, SourceNotFoundError()) + + deliver_trace_event.run({"event": _event().model_dump_json()}) + + +def test_sources_not_ready_are_redelivered_and_keep_their_payload( + monkeypatch: pytest.MonkeyPatch, storage: _Storage +) -> None: + _use_delivery(monkeypatch, 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) -> None: + _use_delivery(monkeypatch, TraceCollectorUnavailableError()) + 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(monkeypatch: pytest.MonkeyPatch, storage: _Storage) -> None: + delivery = _use_delivery(monkeypatch, 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.retried == [job] + assert key not in storage.files + + +def test_unreadable_export_jobs_are_dropped(monkeypatch: pytest.MonkeyPatch, storage: _Storage) -> None: + delivery = _use_delivery(monkeypatch, 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.retried == [] + assert storage.files == {} + + +def test_sweep_removes_expired_and_legacy_payloads_only(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", "other/keep.json"]) diff --git a/api/uv.lock b/api/uv.lock index c950ccddf69f91..c2d447643d61fc 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/docker/envs/core-services/shared.env.example b/docker/envs/core-services/shared.env.example index 1a6280dab03d06..e97c2709ef7daf 100644 --- a/docker/envs/core-services/shared.env.example +++ b/docker/envs/core-services/shared.env.example @@ -111,15 +111,13 @@ 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 +# 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 diff --git a/oxlint-suppressions.json b/oxlint-suppressions.json index 272e853ab21b38..aecb4668853317 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 f82577bbc7c484..8fd71c0dbe1563 100644 --- a/packages/contracts/generated/api/console/apps/errors.gen.ts +++ b/packages/contracts/generated/api/console/apps/errors.gen.ts @@ -787,6 +787,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 a1ff3255300852..0e8f476dee8e2d 100644 --- a/packages/contracts/generated/api/console/apps/orpc.gen.ts +++ b/packages/contracts/generated/api/console/apps/orpc.gen.ts @@ -2827,7 +2827,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 9b7dd626171cf2..f79dc2ff3bb6fb 100644 --- a/packages/contracts/generated/api/console/apps/types.gen.ts +++ b/packages/contracts/generated/api/console/apps/types.gen.ts @@ -944,6 +944,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 @@ -7366,7 +7367,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 5f7e5186894e55..59f4f0bf879629 100644 --- a/packages/contracts/generated/api/console/apps/zod.gen.ts +++ b/packages/contracts/generated/api/console/apps/zod.gen.ts @@ -898,6 +898,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 cbbe2b75614018..4b1ee2767a21eb 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' From ce78971b4b795f84337374b87e7255c0d72772d0 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Tue, 27 Oct 2026 15:37:53 +0800 Subject: [PATCH 02/17] refactor(tracing): put the trace pipeline and Agent trace collection behind ports Align the tracing slice with the layered application-service architecture. - Application services receive ports and plain settings at construction: TraceScopes, TraceDeliveryService, the provider configuration gateway and the Agent trace collector no longer open sessions, read configuration, or reach Celery, storage or the encrypter themselves. The module-level service factories are gone; extensions/application_services/tracing.py builds every tracing service and is reachable as application_services().tracing. - Trace records and Agent collection records are plain values. Repositories map rows to records before their sessions close: trace record assembly moved to repositories/trace_record_repository.py, the Agent trace ledger returns run, invocation and source records, and the tracing configuration repository answers the pipeline's policy and destination queries. Database failures surface as TracingStoreUnavailableError. - Runtime adapters implement the ports: the Celery emitter and retry queue, the entry-point provider catalog, workspace secret encryption, enterprise telemetry and Agent export scheduling. - Tracing errors live in services/tracing/errors.py without HTTP statuses; the Agent tracing handlers map them, and handlers and Celery tasks obtain tracing services from the composition root. - Hooks are bound only when Agent tracing is enabled, so no tracing code reads its configuration outside the composition root. - Import contracts keep the tracing application modules framework and persistence neutral, the tracing handlers off persistence, and the tracing repositories on the contracts they implement. - Tests use in-memory port implementations and real repositories instead of patching module internals. --- api/.importlinter | 76 ++ api/controllers/console/agent/tracing.py | 7 +- api/controllers/console/agent/tracing_runs.py | 7 +- api/controllers/inner_api/agent/tracing.py | 20 +- api/enterprise/telemetry/enterprise_trace.py | 12 +- .../application_services/tracing.py | 93 ++ api/extensions/ext_application_services.py | 34 +- api/extensions/tracing.py | 70 +- api/providers/trace/README.md | 10 +- .../app/tracing_config_repository.py | 53 +- .../application_trace_repository.py | 570 +++++++--- .../trace_record_repository.py} | 79 +- api/repositories/tracing_store_errors.py | 16 + api/services/conversation_name_generator.py | 56 +- .../message_suggested_questions_generator.py | 9 +- api/services/tracing/adapters.py | 114 ++ api/services/tracing/canonical.py | 7 +- api/services/tracing/collection.py | 378 +++++++ api/services/tracing/delivery.py | 169 ++- api/services/tracing/emitter.py | 129 --- api/services/tracing/errors.py | 59 ++ api/services/tracing/provider_configs.py | 159 ++- api/services/tracing/providers.py | 104 ++ api/services/tracing/records.py | 75 +- api/services/tracing/recovery.py | 217 ++-- api/services/tracing/registry.py | 56 +- api/services/tracing/runtime_bridge.py | 245 ++--- api/services/tracing/scopes.py | 77 ++ api/services/tracing/service.py | 482 +++++---- api/tasks/application_trace_task.py | 13 +- api/tasks/trace_delivery_task.py | 13 +- ...st_message_suggested_questions_failures.py | 5 +- .../app/test_message_suggested_questions.py | 8 +- .../explore/test_trial_suggested_questions.py | 3 +- .../app/test_conversation_rename.py | 5 +- .../app/test_message_suggested_questions.py | 5 +- .../web/test_conversation_rename.py | 5 +- .../web/test_message_suggested_questions.py | 7 +- .../helpers/unit/services/tracing_ports.py | 175 ++++ api/tests/unit_tests/conftest.py | 4 +- .../controllers/console/agent/test_tracing.py | 17 +- .../console/agent/test_tracing_runs.py | 103 +- .../controllers/console/app/test_ops_trace.py | 4 +- .../inner_api/agent/test_tracing.py | 167 ++- .../agent_app/test_runtime_request_builder.py | 26 +- .../agent_v2/test_runtime_request_builder.py | 30 +- .../enterprise/telemetry/test_draft_trace.py | 6 +- .../telemetry/test_enterprise_trace.py | 43 +- .../application_services/test_tracing.py | 81 ++ .../test_ext_application_services.py | 38 +- .../unit_tests/extensions/test_tracing.py | 145 ++- .../app/test_tracing_config_repository.py | 88 +- .../test_application_trace_repository.py | 389 +++++-- .../test_trace_record_repository.py} | 40 +- .../knowledge_fs/test_workflow_runtime.py | 7 +- .../test_agent_tracing_config_service.py | 4 +- .../test_app_tracing_config_service.py | 3 +- .../test_conversation_name_generator.py | 67 +- ...t_message_suggested_questions_generator.py | 28 +- ...est_message_suggested_questions_service.py | 5 +- .../services/test_workflow_persistence.py | 68 +- .../services/tracing/test_adapters.py | 205 ++++ .../services/tracing/test_canonical.py | 20 +- .../services/tracing/test_delivery.py | 334 +++--- .../services/tracing/test_emitter.py | 202 ---- .../services/tracing/test_provider_configs.py | 146 +-- .../services/tracing/test_providers.py | 90 ++ .../services/tracing/test_recovery.py | 570 +++++++--- .../services/tracing/test_registry.py | 25 +- .../services/tracing/test_runtime_bridge.py | 988 +++++++++++------- .../services/tracing/test_scopes.py | 103 ++ .../services/tracing/test_service.py | 908 +++++++++------- .../tasks/test_application_trace_task.py | 135 ++- .../tasks/test_trace_delivery_task.py | 176 +++- 74 files changed, 5897 insertions(+), 2990 deletions(-) create mode 100644 api/extensions/application_services/tracing.py rename api/{services/tracing/assembly.py => repositories/trace_record_repository.py} (93%) create mode 100644 api/repositories/tracing_store_errors.py create mode 100644 api/services/tracing/adapters.py create mode 100644 api/services/tracing/collection.py delete mode 100644 api/services/tracing/emitter.py create mode 100644 api/services/tracing/errors.py create mode 100644 api/services/tracing/providers.py create mode 100644 api/services/tracing/scopes.py create mode 100644 api/tests/helpers/unit/services/tracing_ports.py create mode 100644 api/tests/unit_tests/extensions/application_services/test_tracing.py rename api/tests/unit_tests/{services/tracing/test_assembly.py => repositories/test_trace_record_repository.py} (81%) create mode 100644 api/tests/unit_tests/services/tracing/test_adapters.py delete mode 100644 api/tests/unit_tests/services/tracing/test_emitter.py create mode 100644 api/tests/unit_tests/services/tracing/test_providers.py create mode 100644 api/tests/unit_tests/services/tracing/test_scopes.py diff --git a/api/.importlinter b/api/.importlinter index 6f7a91ab4467aa..381631cf3f9347 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -1299,3 +1299,79 @@ 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.delivery + services.tracing.errors + services.tracing.hierarchy + services.tracing.normalization + services.tracing.provider_configs + services.tracing.providers + services.tracing.records + services.tracing.recovery + services.tracing.scopes + services.tracing.service +forbidden_modules = + celery + configs + controllers + 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 +# 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/controllers/console/agent/tracing.py b/api/controllers/console/agent/tracing.py index 3cbb97ebf2b1d7..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 @@ -49,7 +49,6 @@ AppTracingConfigProviderUnavailableError, AppTracingConfigVerificationFailedError, ) -from services.tracing.provider_configs import capabilities_payload _EDIT_ROLES = frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR}) @@ -209,14 +208,14 @@ def delete(self, context: RequestContext, agent_id: UUID): def _provider_capabilities(): query = TraceProviderQuery.model_validate(request.args.to_dict(flat=True)) try: - capabilities = capabilities_payload(query.tracing_provider) + 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}, ) 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/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/enterprise/telemetry/enterprise_trace.py b/api/enterprise/telemetry/enterprise_trace.py index 5a9f27b5fcd8c9..020feaf1ceb52e 100644 --- a/api/enterprise/telemetry/enterprise_trace.py +++ b/api/enterprise/telemetry/enterprise_trace.py @@ -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/extensions/application_services/tracing.py b/api/extensions/application_services/tracing.py new file mode 100644 index 00000000000000..7944f5c90dc9a9 --- /dev/null +++ b/api/extensions/application_services/tracing.py @@ -0,0 +1,93 @@ +"""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.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() + 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_traces = ApplicationTracingService( + traces=ApplicationTraceRepository(database_client), + # App, workspace, and credential names are enterprise telemetry dimensions only. + records=TraceRecordRepository(database_client, 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, 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 145c46ea9057c8..e960a2e9b6f85f 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -20,7 +20,6 @@ from core.schemas.schema_manager import SchemaManager from core.tools.tool_file_manager import ToolFileManager from core.tools.tool_manager import ToolManager -from enterprise.telemetry.exporter import is_enterprise_telemetry_enabled from enums import DeploymentEdition, WebAppAccessMode from extensions.application_services.account import AccountServices, build_account_services from extensions.application_services.agent import AgentAppServices, build_agent_app_services @@ -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, @@ -138,7 +138,7 @@ 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,10 +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.egress import SSRFEgress -from services.tracing.emitter import CeleryTraceEmitter, TraceScopes -from services.tracing.provider_configs import ProviderConfigGateway -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 ( @@ -369,6 +366,7 @@ class ApplicationServices: web_passport: WebPassportService tags: TagApplicationService workflow_statistics: WorkflowStatisticQueryService + tracing: TracingServices def _build_oauth_server_service( @@ -519,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=ProviderConfigGateway(SSRFEgress()), + provider=tracing.provider_configs, ) apps = build_app_services( database_client=database_client, @@ -587,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 ) @@ -772,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), @@ -869,6 +871,7 @@ def build_application_services( session_maker=database_client, ), ), + tracing=tracing, ) @@ -877,23 +880,14 @@ def init_app(app: Flask) -> None: app.extensions["knowledge_fs_node_dependencies"] = build_knowledge_fs_node_dependencies ToolManager.register_hardcoded_provider_class(KnowledgeFSProvider) - session_maker = get_session_maker() services = build_application_services( - database_client=session_maker, + database_client=get_session_maker(), deployment_edition=dify_config.DEPLOYMENT_EDITION, initialization_password=dify_config.INIT_PASSWORD, redis=redis_client, ) app.extensions[_EXTENSION_KEY] = services - bind_tracing( - app, - scopes=TraceScopes( - emitter=CeleryTraceEmitter(), - session_factory=session_maker, - enterprise_enabled=is_enterprise_telemetry_enabled, - ), - hooks=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/tracing.py b/api/extensions/tracing.py index cb33ae906fc2cb..93e60916e43063 100644 --- a/api/extensions/tracing.py +++ b/api/extensions/tracing.py @@ -1,8 +1,8 @@ -"""Execution-side ports of the trace pipeline, bound to one Flask application. +"""Execution-side access to the trace pipeline, bound to one Flask application. -Core producers import only these accessors; the composition root binds the -service implementations. Every accessor fails open: tracing never breaks the -operation that is being traced. +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 @@ -10,9 +10,10 @@ from typing import TYPE_CHECKING, Any, Literal, Protocol, cast from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation from flask import Flask, current_app, has_app_context -from configs import dify_config +from core.helper.trace_id_helper import ParentTraceContext if TYPE_CHECKING: from dify_agent.layers.execution_context import ApplicationTracingContext @@ -27,19 +28,21 @@ _HOOKS_KEY = "application_tracing_hooks" -class TraceScopeFactory(Protocol): - def __call__( +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, - extras: Mapping[str, Any], + correlation: TraceCorrelation, ) -> TraceScope | None: ... def for_app( - self, *, tenant_id: str, app_id: str | None, user_id: str | None, extras: Mapping[str, Any] + self, *, tenant_id: str, app_id: str | None, user_id: str | None, correlation: TraceCorrelation ) -> TraceScope | None: ... @@ -73,15 +76,33 @@ def capture_agent_app_terminal( ) -> None: ... -def bind_tracing(app: Flask, *, scopes: TraceScopeFactory, hooks: ApplicationTracingHooks) -> 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() -> TraceScopeFactory | None: +def _scopes() -> ExecutionTraceScopes | None: if not has_app_context(): return None - return cast(TraceScopeFactory | None, current_app.extensions.get(_SCOPES_KEY)) + return cast(ExecutionTraceScopes | None, current_app.extensions.get(_SCOPES_KEY)) + + +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( @@ -92,7 +113,13 @@ def create_trace_scope( if scopes is None: return None try: - return scopes(tenant_id=tenant_id, app_id=app_id, app_tracing=app_tracing, user_id=user_id, extras=extras) + 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 @@ -106,18 +133,19 @@ def create_trace_scope_for_app( if scopes is None: return None try: - return scopes.for_app(tenant_id=tenant_id, app_id=app_id, user_id=user_id, extras=extras) + 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, or ``None`` when Agent tracing is off or cannot run here.""" - if not dify_config.OPS_AGENT_TRACE_ENABLED: + """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 = cast(ApplicationTracingHooks | None, current_app.extensions.get(_HOOKS_KEY)) if has_app_context() else None - if hooks is None: - # Agent runs started here are reported as not collected instead of failing. - logger.warning("Agent tracing is enabled but its hooks are not bound in this context") - return hooks + return cast(ApplicationTracingHooks | None, current_app.extensions.get(_HOOKS_KEY)) diff --git a/api/providers/trace/README.md b/api/providers/trace/README.md index e896040bc739c5..35749e6b6e2519 100644 --- a/api/providers/trace/README.md +++ b/api/providers/trace/README.md @@ -71,7 +71,7 @@ Stored configurations keep the provider's fields at the top level, with `secret( = "dify_trace_.provider:Provider" ``` -3. Add the ID to `SUPPORTED_PROVIDERS` in `api/services/tracing/registry.py` and the provider's form fields to the web console's tracing field registry. +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. @@ -84,9 +84,11 @@ Stored configurations keep the provider's fields at the top level, with `secret( - moderation, tools, retrieval, conversation naming and suggested questions. A scope exists only when the app has tracing enabled or enterprise telemetry is on. -2. `tasks.trace_delivery_task` delivers each event on the `ops_trace` queue: - - `services.tracing.assembly` reloads the records through the event's workspace. +2. `tasks.trace_delivery_task` hands each event on the `ops_trace` queue 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 takes the executions it manages. + - 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. 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 689b2a2e91a976..f9b4ab74760245 100644 --- a/api/repositories/application_trace_repository.py +++ b/api/repositories/application_trace_repository.py @@ -1,13 +1,21 @@ -"""Bounded transactions for application trace collection and delivery leases.""" +"""Bounded transactions for Agent trace collection and delivery leases. + +Rows become collection 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 typing import Any, override from uuid import UUID, uuid5 from dify_agent.protocol.tracing import AgentTraceBatch from dify_trace_sdk.ids import span_id_for +from dify_trace_sdk.provider import SendResult from sqlalchemy import delete, or_, select from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session, sessionmaker @@ -15,62 +23,285 @@ 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 repositories.tracing_store_errors import tracing_store_errors +from services.tracing.collection import ( + CLOSED_STATES, + ApplicationTraceStore, + MessageProgress, + SourceMessage, + SourceNodeExecution, + SourceWorkflowRun, + TraceApp, + TraceAppMode, + TraceBatch, + TraceInvocation, + TraceRun, + TraceRunState, + TraceSourceKind, +) +from services.tracing.errors import ( + TraceCollectionConflictError, + TraceCollectionLimitError, + TraceCollectionNotFoundError, +) from services.tracing.normalization import safe_attributes +_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 _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, + ) + -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(ApplicationTraceStore): def __init__(self, sessions: sessionmaker[Session]) -> None: - self.sessions = sessions + self._sessions = sessions - def _source( - self, *, tenant_id: str, app_id: str, kind: str, source_id: str + @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 _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 + @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 + 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 + ) + ) + + @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 + 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) + ) + + @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, + ) + ).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 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: + workflow_id: str | None, + ) -> TraceRun | 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: + with self._read() as session: + if self._owned_app_id(session, tenant_id=tenant_id, app_id=app_id) is None: return None existing = session.scalar( select(ApplicationTraceRun).where( @@ -81,9 +312,9 @@ def start( ) ) 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) - with self.sessions.begin() as session: + return _run(existing) + source = self._source_row(tenant_id=tenant_id, app_id=app_id, kind=source_kind, source_id=source_id) + with self._write() as session: app = session.scalar(select(App).where(App.id == app_id, App.tenant_id == tenant_id)) if app is None: return None @@ -96,7 +327,7 @@ def start( ) ) if existing is not None: - return existing + return _run(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 @@ -148,15 +379,16 @@ def start( 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", + state=TraceRunState.COLLECTING if enabled else TraceRunState.NOT_COLLECTED, ) - return self._insert_or_get(session, record) + return _run(self._insert_or_get(session, record)) + @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, @@ -164,22 +396,33 @@ def source_run( ApplicationTraceRun.source_id == source_id, ) ) + return _run(run) if run is not None else None + @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) + .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, + ) + ) + return _run(run) if run is not None else None + + @override def record_initialization_failure( - self, - *, - tenant_id: str, - app_id: str, - source_kind: str, - source_id: str, - ) -> ApplicationTraceRun | None: + self, *, tenant_id: str, app_id: str, source_kind: TraceSourceKind, source_id: str + ) -> TraceRun | 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) + source = self._source_row(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)) - if app is None: + with self._write() as session: + if self._owned_app_id(session, tenant_id=tenant_id, app_id=app_id) is None: return None current = session.scalar( select(ApplicationTraceRun).where( @@ -190,7 +433,7 @@ def record_initialization_failure( ) ) if current is not None: - return current + return _run(current) current = ApplicationTraceRun( id=str(uuidv7()), tenant_id=tenant_id, @@ -205,11 +448,11 @@ def record_initialization_failure( session_id=source_id, capture_content=False, capture_agent=False, - state="failed", + state=TraceRunState.FAILED, error_code="tracing_initialization_failed", terminal_at=naive_utc_now(), ) - return self._insert_or_get(session, current) + return _run(self._insert_or_get(session, current)) @staticmethod def _insert_or_get(session: Session, record: ApplicationTraceRun) -> ApplicationTraceRun: @@ -237,21 +480,26 @@ def _insert_or_get(session: Session, record: ApplicationTraceRun) -> Application return existing return record + @staticmethod + def _locked_run(session: Session, run_id: str) -> ApplicationTraceRun | None: + return session.scalar(select(ApplicationTraceRun).where(ApplicationTraceRun.id == run_id).with_for_update()) + + @override def invocation( self, *, run_id: str, agent_id: str, node_execution_id: str | None, attempt: int, call_key: str - ) -> ApplicationTraceInvocation: + ) -> TraceInvocation: 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) + with self._write() as session: + run = self._locked_run(session, run_id) + if run is None or run.state not in {TraceRunState.COLLECTING, TraceRunState.READY}: + raise TraceCollectionConflictError("tracing_run_not_collecting") 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) + raise TraceCollectionConflictError("tracing_invocation_conflict") if invocation is None: invocation = ApplicationTraceInvocation( id=identifier, @@ -263,23 +511,21 @@ def invocation( attempt=attempt, ) session.add(invocation) - return invocation + return _invocation(invocation) + @override 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: - run = session.scalar( - select(ApplicationTraceRun).where(ApplicationTraceRun.id == str(batch.run_id)).with_for_update() - ) + with self._write() as session: + run = self._locked_run(session, str(batch.run_id)) 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) + raise TraceCollectionNotFoundError("tracing_run_not_found") + if self._owned_app_id(session, tenant_id=run.tenant_id, app_id=run.app_id) is None: + raise TraceCollectionNotFoundError("tracing_owner_not_found") invocation = session.scalar( select(ApplicationTraceInvocation).where( ApplicationTraceInvocation.id == str(batch.invocation_id), @@ -287,7 +533,7 @@ def ingest(self, batch: AgentTraceBatch, *, max_spans: int, max_bytes: int) -> t ) ) if invocation is None: - raise TraceCollectionError("tracing_invocation_not_found", 404) + raise TraceCollectionNotFoundError("tracing_invocation_not_found") previous = session.get(ApplicationTraceBatch, str(batch.batch_id)) if previous is not None: if ( @@ -295,14 +541,14 @@ def ingest(self, batch: AgentTraceBatch, *, max_spans: int, max_bytes: int) -> t or previous.invocation_id != invocation.id or previous.payload_hash != digest ): - raise TraceCollectionError("tracing_batch_conflict", 409) + raise TraceCollectionConflictError("tracing_batch_conflict") return len(previous.spans), run.id - if run.state in {"accepted", "exporting", "uncertain", "failed"} or invocation.finalized: - raise TraceCollectionError("tracing_collection_closed", 409) + if run.state in CLOSED_STATES or invocation.finalized: + raise TraceCollectionConflictError("tracing_collection_closed") if any(span.trace_id != run.trace_id for span in batch.spans): - raise TraceCollectionError("tracing_trace_id_mismatch", 409) + raise TraceCollectionConflictError("tracing_trace_id_mismatch") if run.span_count + len(batch.spans) > max_spans or run.payload_bytes + len(canonical) > max_bytes: - raise TraceCollectionError("tracing_payload_limit", 413) + raise TraceCollectionLimitError("tracing_payload_limit") spans: list[dict[str, Any]] = [] for value in batch.spans: span = value.model_dump(mode="json") @@ -329,60 +575,76 @@ def ingest(self, batch: AgentTraceBatch, *, max_spans: int, max_bytes: int) -> t 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"}: + @override + def terminal(self, run_id: str, trace: Mapping[str, Any], root_span_id: str) -> None: + with self._write() as session: + run = self._locked_run(session, run_id) + if run is None or run.state in CLOSED_STATES: 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.outer_trace = dict(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) + @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 - 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, - ) + @staticmethod + def _invocation_rows(session: Session, run_id: str) -> list[TraceInvocation]: + return [ + _invocation(row) + for row in session.scalars( + select(ApplicationTraceInvocation) + .where(ApplicationTraceInvocation.run_id == run_id) + .order_by( + ApplicationTraceInvocation.attempt, + ApplicationTraceInvocation.created_at, + ApplicationTraceInvocation.id, ) ) - batches = list( - session.scalars( + ] + + @override + def invocations(self, run_id: str) -> list[TraceInvocation]: + with self._read() as session: + return self._invocation_rows(session, run_id) + + @override + def collected(self, run_id: str) -> tuple[list[TraceInvocation], list[TraceBatch]]: + with self._read() as session: + batches = [ + TraceBatch(invocation_id=row.invocation_id, spans=list(row.spans)) + for row in session.scalars( select(ApplicationTraceBatch) .where(ApplicationTraceBatch.run_id == run_id) .order_by(ApplicationTraceBatch.created_at) ) - ) - return invocations, batches + ] + return self._invocation_rows(session, run_id), batches - def claim( - self, run_id: str, *, collection_timeout: int, lease_seconds: int - ) -> tuple[ApplicationTraceRun, str] | None: + @override + def claim(self, run_id: str, *, collection_timeout: int, lease_seconds: int) -> tuple[TraceRun, 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"}: + with self._write() as session: + run = self._locked_run(session, run_id) + if run is None or run.state in { + TraceRunState.NOT_COLLECTED, + TraceRunState.ACCEPTED, + TraceRunState.FAILED, + TraceRunState.UNCERTAIN, + }: return None - if run.state == "exporting": + if run.state == TraceRunState.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.state, run.error_code = TraceRunState.UNCERTAIN, "export_lease_expired" run.updated_at = now return None if run.outer_trace is None: @@ -394,56 +656,57 @@ def claim( ) 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.state, run.error_code = TraceRunState.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" + if self._owned_app_id(session, tenant_id=run.tenant_id, app_id=run.app_id) is None: + run.state, run.error_code = TraceRunState.FAILED, "tracing_owner_not_found" return None token = str(uuidv7()) - run.state = "exporting" + run.state = TraceRunState.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 + return _run(run), token + @override 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: + with self._write() as session: + run = self._locked_run(session, run_id) + if run is None or run.state != TraceRunState.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 + @override def complete( self, run_id: str, token: str, *, - state: str, - error_code: str | None = None, - result: Any = None, - retry_delay: int | None = None, + state: TraceRunState, + error_code: str | None, + result: SendResult | None, + retry_delay: int | None, ) -> None: - with self.sessions.begin() as session: - run = session.scalar(select(ApplicationTraceRun).where(ApplicationTraceRun.id == run_id).with_for_update()) + with self._write() as session: + run = self._locked_run(session, run_id) if run is None or run.export_token != token: return late_acceptance = ( - state == "accepted" + state is TraceRunState.ACCEPTED and result is not None - and run.state == "uncertain" + and run.state == TraceRunState.UNCERTAIN and run.error_code == "export_lease_expired" ) - if run.state != "exporting" and not late_acceptance: + if run.state != TraceRunState.EXPORTING and not late_acceptance: return run.state, run.error_code = state, error_code run.updated_at = naive_utc_now() @@ -452,28 +715,41 @@ def complete( run.provider_trace_id, run.trace_url, run.span_urls = ( result.trace_id, result.trace_url, - result.span_urls, + dict(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: + @override + def recovery_checked(self, run_id: str, *, error_code: str | 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"}: + with self._write() as session: + run = self._locked_run(session, run_id) + if ( + run is None + or run.outer_trace is not None + or run.state not in {TraceRunState.COLLECTING, TraceRunState.READY} + ): return run.updated_at = naive_utc_now() if error_code is not None: - run.state, run.error_code = "failed", error_code + run.state, run.error_code = TraceRunState.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(), @@ -484,15 +760,16 @@ def pending_ids(self, limit: int = 100) -> list[str]: ) ) + @override 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) + with self._write() as session: + run = self._locked_run(session, run_id) + if run is None or run.state == TraceRunState.NOT_COLLECTED: + raise TraceCollectionNotFoundError("tracing_run_not_found") + if run.state == TraceRunState.ACCEPTED: + raise TraceCollectionConflictError("tracing_already_accepted") + if run.state not in {TraceRunState.FAILED, TraceRunState.UNCERTAIN} or run.outer_trace is None: + raise TraceCollectionConflictError("tracing_retry_not_allowed") incomplete = session.scalar( select(ApplicationTraceInvocation.id) .where( @@ -502,16 +779,17 @@ def retry(self, *, run_id: str, acknowledge_possible_duplicates: bool) -> None: .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 + raise TraceCollectionConflictError("tracing_retry_not_allowed") + if run.state == TraceRunState.UNCERTAIN and not acknowledge_possible_duplicates: + raise TraceCollectionConflictError("tracing_duplicate_acknowledgement_required") + run.state, run.error_code, run.export_token = TraceRunState.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/services/tracing/assembly.py b/api/repositories/trace_record_repository.py similarity index 93% rename from api/services/tracing/assembly.py rename to api/repositories/trace_record_repository.py index 1d42426f3c6b93..0b6bd57a1f3522 100644 --- a/api/services/tracing/assembly.py +++ b/api/repositories/trace_record_repository.py @@ -2,13 +2,14 @@ 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. +and facts persisted nowhere else, so records reflect committed state. Rows +become plain records before their session closes. """ import json from collections.abc import Mapping, Sequence from datetime import datetime, timedelta -from typing import Any +from typing import Any, override from dify_trace_sdk.events import ( ConversationNamed, @@ -35,6 +36,9 @@ from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider from models.workflow import WorkflowAppLog, WorkflowNodeExecutionTriggeredFrom, WorkflowRun from repositories.factory import DifyAPIRepositoryFactory +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, @@ -45,6 +49,7 @@ PromptGenerationTraceInfo, SuggestedQuestionTraceInfo, ToolTraceInfo, + TraceMessage, WorkflowTraceInfo, ) @@ -67,14 +72,6 @@ } -class SourceNotFoundError(Exception): - """The referenced record does not exist for the event's tenant and app.""" - - -class SourceNotReadyError(Exception): - """The referenced record is not persisted or not finished yet.""" - - 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"): @@ -91,6 +88,23 @@ def _text_input(inputs: Mapping[str, Any], *names: str) -> str: 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 {} @@ -109,15 +123,15 @@ def _streaming_metrics(message: Message) -> dict[str, Any]: } -class TraceAssembler: +class TraceRecordRepository(TraceRecordSource): def __init__(self, session_factory: sessionmaker[Session], *, files_url: str, enterprise_details: bool) -> None: self._sessions = session_factory 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: - """Raises SourceNotFoundError or SourceNotReadyError.""" match event.payload: case MessageCompleted() as payload: return self._message(event, payload) @@ -211,7 +225,7 @@ def _correlation_metadata(event: TraceEvent) -> dict[str, Any]: # Records def _message(self, event: TraceEvent, payload: MessageCompleted) -> MessageTraceInfo: - with self._sessions() as session: + 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( @@ -241,7 +255,7 @@ def _message(self, event: TraceEvent, payload: MessageCompleted) -> MessageTrace return MessageTraceInfo( trace_id=event.correlation.external_trace_id, message_id=message.id, - message_data=message, + message_data=_trace_message(message), conversation_model=conversation_mode, message_tokens=message.message_tokens, answer_tokens=message.answer_tokens, @@ -265,13 +279,14 @@ def _workflow(self, event: TraceEvent, payload: WorkflowRunCompleted) -> Workflo if event.app_id is None: raise SourceNotFoundError runs = DifyAPIRepositoryFactory.create_api_workflow_run_repository(self._sessions) - run = runs.get_workflow_run_by_id( - tenant_id=event.tenant_id, app_id=event.app_id, run_id=payload.workflow_run_id - ) + with tracing_store_errors(): + run = 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 self._sessions() as session: + with tracing_store_errors(), self._sessions() as session: self._app(session, event) workflow_app_log_id = session.scalar( select(WorkflowAppLog.id).where( @@ -322,7 +337,6 @@ def _workflow(self, event: TraceEvent, payload: WorkflowRunCompleted) -> Workflo } return WorkflowTraceInfo( trace_id=correlation.external_trace_id, - workflow_data=run, node_executions=executions, conversation_id=payload.conversation_id, message_started_at=message_row[1] if message_row else None, @@ -360,7 +374,8 @@ def _node_executions(self, run: WorkflowRun) -> Sequence[WorkflowNodeExecution]: app_id=run.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ) - return repository.get_by_workflow_execution(workflow_execution_id=run.id) + with tracing_store_errors(): + return repository.get_by_workflow_execution(workflow_execution_id=run.id) @staticmethod def _token_split(executions: Sequence[WorkflowNodeExecution]) -> tuple[int, int]: @@ -378,7 +393,7 @@ def _token_split(executions: Sequence[WorkflowNodeExecution]) -> tuple[int, int] return prompt_total, completion_total def _moderation(self, event: TraceEvent, payload: ModerationChecked) -> ModerationTraceInfo: - with self._sessions() as session: + 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) @@ -396,7 +411,7 @@ def _moderation(self, event: TraceEvent, payload: ModerationChecked) -> Moderati # Enterprise telemetry keys chatflow moderation by its workflow log. message_id=workflow_app_log_id or message.id, inputs=payload.inputs, - message_data=message, + message_data=_trace_message(message), flagged=payload.flagged, action=payload.action, preset_response=payload.preset_response, @@ -409,7 +424,7 @@ def _moderation(self, event: TraceEvent, payload: ModerationChecked) -> Moderati def _suggested_questions( self, event: TraceEvent, payload: SuggestedQuestionsGenerated ) -> SuggestedQuestionTraceInfo: - with self._sessions() as session: + 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) @@ -422,7 +437,7 @@ def _suggested_questions( return SuggestedQuestionTraceInfo( trace_id=event.correlation.external_trace_id, message_id=workflow_app_log_id or message.id, - message_data=message, + message_data=_trace_message(message), inputs=message.message, outputs=message.answer, start_time=payload.started_at, @@ -448,7 +463,7 @@ def _dataset_retrieval(self, event: TraceEvent, payload: DatasetRetrieved) -> Da 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 self._sessions() as session: + 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) @@ -485,12 +500,12 @@ def _dataset_retrieval(self, event: TraceEvent, payload: DatasetRetrieved) -> Da start_time=payload.started_at, end_time=payload.ended_at, metadata=metadata, - message_data=message, + message_data=_trace_message(message), error=None, ) def _tool(self, event: TraceEvent, payload: ToolInvoked) -> ToolTraceInfo: - with self._sessions() as session: + 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) @@ -524,7 +539,7 @@ def _tool(self, event: TraceEvent, payload: ToolInvoked) -> ToolTraceInfo: return ToolTraceInfo( trace_id=event.correlation.external_trace_id, message_id=message.id, - message_data=message, + message_data=_trace_message(message), tool_name=payload.tool_name, start_time=payload.started_at, end_time=payload.ended_at, @@ -541,7 +556,7 @@ def _tool(self, event: TraceEvent, payload: ToolInvoked) -> ToolTraceInfo: ) def _conversation_name(self, event: TraceEvent, payload: ConversationNamed) -> GenerateNameTraceInfo: - with self._sessions() as session: + with tracing_store_errors(), self._sessions() as session: if event.app_id is not None: self._app(session, event) message = ( @@ -557,7 +572,7 @@ def _conversation_name(self, event: TraceEvent, payload: ConversationNamed) -> G return GenerateNameTraceInfo( trace_id=event.correlation.external_trace_id, message_id=payload.message_id, - message_data=message, + message_data=_trace_message(message) if message is not None else None, conversation_id=payload.conversation_id, inputs=payload.inputs, outputs=payload.name, @@ -568,7 +583,7 @@ def _conversation_name(self, event: TraceEvent, payload: ConversationNamed) -> G ) def _prompt_generation(self, event: TraceEvent, payload: PromptGenerated) -> PromptGenerationTraceInfo: - with self._sessions() as session: + 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) @@ -611,7 +626,7 @@ def _draft_node(self, event: TraceEvent, payload: DraftNodeExecuted) -> DraftNod node = payload.node_execution if node.get("tenant_id") != event.tenant_id or node.get("app_id") != event.app_id: raise SourceNotFoundError - with self._sessions() as session: + 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] = { 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/services/conversation_name_generator.py b/api/services/conversation_name_generator.py index 7a8a86f9475f21..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 @@ -8,28 +10,38 @@ from core.model_context import use_credit_usage_metadata from core.model_manager import ModelManager from extensions.ext_database import db -from extensions.tracing import create_trace_scope_for_app 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, - trace_scope=create_trace_scope_for_app(tenant_id=tenant_id, app_id=app_id, user_id=None, extras={}), - ) + @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 bc791e72f9ee28..ca2036097b5873 100644 --- a/api/services/message_suggested_questions_generator.py +++ b/api/services/message_suggested_questions_generator.py @@ -17,12 +17,12 @@ 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 extensions.tracing import create_trace_scope_for_app 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, @@ -131,8 +134,8 @@ def generate( # Invocation catches provider failures, including failed database # transactions. Tracing always starts with a fresh scoped session. db.session.remove() - trace_scope = create_trace_scope_for_app( - tenant_id=context.tenant_id, app_id=context.app_id, user_id=None, extras={} + 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( diff --git a/api/services/tracing/adapters.py b/api/services/tracing/adapters.py new file mode 100644 index 00000000000000..802454e55cf207 --- /dev/null +++ b/api/services/tracing/adapters.py @@ -0,0 +1,114 @@ +"""Runtime adapters of the trace pipeline over Celery, storage, workspace keys, and enterprise telemetry. + +Events go straight to the delivery worker; large ones travel through storage +instead of the broker. Every payload key carries its date so the sweeper can +expire payloads that a crashed delivery left behind. +""" + +from collections.abc import Sequence +from datetime import UTC, datetime +from typing import override +from uuid import uuid4 + +from celery import current_app as celery_app +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" +TRACE_QUEUE = "ops_trace" +PAYLOAD_PREFIX = "trace_payloads" +INLINE_EVENT_BYTES = 256 * 1024 +# Node executions may be saved by another worker after their run row; give them time to land. +WORKFLOW_SETTLE_SECONDS = 5 + + +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): + @override + def emit(self, event: TraceEvent) -> None: + encoded = event.model_dump_json() + if len(encoded.encode()) > INLINE_EVENT_BYTES: + key = payload_key(event.occurred_at, "events", event.event_id) + storage.save(key, encoded.encode()) + envelope = {"payload_key": key} + else: + envelope = {"event": encoded} + countdown = WORKFLOW_SETTLE_SECONDS if event.payload.kind == "workflow_run" else 0 + celery_app.send_task(DELIVER_TASK, args=(envelope,), queue=TRACE_QUEUE, countdown=countdown) + + +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,), queue=TRACE_QUEUE, 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,), queue=TRACE_QUEUE) + + +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 index 0283af7a91885d..505fb01a7a645b 100644 --- a/api/services/tracing/canonical.py +++ b/api/services/tracing/canonical.py @@ -220,11 +220,10 @@ def _message(self, record: MessageTraceInfo) -> Trace: ) def _workflow(self, record: WorkflowTraceInfo) -> Trace: - run = record.workflow_data executions = record.node_executions hierarchy = build_workflow_hierarchy(executions) - workflow_start = run.created_at - workflow_end = max(run.finished_at or record.end_time or workflow_start, workflow_start) + 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] = [] @@ -243,7 +242,7 @@ def _workflow(self, record: WorkflowTraceInfo) -> Trace: name="message", kind=SpanKind.CHAIN, start_time=message_start, - end_time=max(record.end_time or workflow_end, 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)), diff --git a/api/services/tracing/collection.py b/api/services/tracing/collection.py new file mode 100644 index 00000000000000..ebb8a71ea08561 --- /dev/null +++ b/api/services/tracing/collection.py @@ -0,0 +1,378 @@ +"""Contracts of Agent trace collection: run records, the ledger port, 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. + +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. +""" + +import json +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from enum import StrEnum +from typing import Any, Literal, Protocol + +from dify_agent.protocol.tracing import AgentTraceBatch +from dify_trace_sdk.provider import SendResult + +from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data +from services.tracing.provider_configs import TraceProviderSupport + +type TraceSourceKind = Literal["message", "workflow", "node"] + + +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" + + +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.""" + + +# A closed run accepts no runtime batches and no other business terminal. +CLOSED_STATES = frozenset( + {TraceRunState.ACCEPTED, TraceRunState.EXPORTING, TraceRunState.UNCERTAIN, TraceRunState.FAILED} +) + + +@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 TraceInvocation: + 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 TraceBatch: + invocation_id: str + spans: Sequence[Mapping[str, Any]] + + +@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.""" + + +@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 + + +type TraceSource = SourceMessage | SourceWorkflowRun | SourceNodeExecution + + +@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 + + +@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 + + +type ScheduleTraceExport = Callable[[str], object] + + +class ApplicationTraceStore(Protocol): + """The Agent trace ledger. Reads and writes are scoped by the tenant and app that own the run. + + Every operation raises TracingStoreUnavailableError when the database + cannot be reached, and the TraceCollectionError subtypes it names. + """ + + def app(self, *, tenant_id: str, app_id: str) -> TraceApp | None: ... + + 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: ... + + def start( + self, + *, + tenant_id: str, + app_id: str, + source_kind: TraceSourceKind, + source_id: str, + session_id: str, + workflow_id: str | None, + ) -> TraceRun | None: + """Freeze the app's destination and capture policy for a source, once. + + A run collects only when the app's enabled provider has a stored + configuration; otherwise it is frozen as not collected. The unique source + key keeps the first frozen run when executions race. ``workflow_id`` lets + a trusted workflow builder reserve a run before its source row is saved. + ``None`` when the source or its owner does not exist. + """ + ... + + 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 record_initialization_failure( + 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.""" + ... + + def invocation( + self, *, run_id: str, agent_id: str, node_execution_id: str | None, attempt: int, call_key: str + ) -> TraceInvocation: + """Create the invocation of a call key once, under the run's row lock. + + Raises: + TraceCollectionConflictError: The run stopped collecting, or the call key names another invocation. + """ + ... + + def ingest(self, batch: AgentTraceBatch, *, max_spans: int, max_bytes: int) -> tuple[int, str]: + """Append a runtime batch under the run's row lock and return its span count and run ID. + + A batch ID replays only with the same payload hash; a final batch finalizes its invocation. + + Raises: + TraceCollectionNotFoundError: The run, its owner, or the invocation does not exist. + TraceCollectionConflictError: The batch conflicts with the ledger or the run stopped collecting. + TraceCollectionLimitError: The batch exceeds the run's limits. + """ + ... + + def terminal(self, run_id: str, trace: Mapping[str, Any], root_span_id: str) -> None: + """Persist the business terminal once under the run's row lock; a different later terminal is ignored.""" + ... + + def get(self, run_id: str) -> TraceRun | None: ... + + def invocations(self, run_id: str) -> list[TraceInvocation]: + """The run's invocations by attempt, then creation.""" + ... + + def collected(self, run_id: str) -> tuple[list[TraceInvocation], list[TraceBatch]]: + """The run's invocations by attempt, then creation, and its batches by arrival.""" + ... + + def claim(self, run_id: str, *, collection_timeout: int, lease_seconds: int) -> tuple[TraceRun, str] | None: + """Lease a complete run for one export attempt and return it with the attempt's token. + + 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. The run's row lock + serializes claims; the returned token fences the attempt. + """ + ... + + def renew_export_lease(self, run_id: str, token: str) -> bool: ... + + def complete( + self, + run_id: str, + token: str, + *, + state: TraceRunState, + error_code: str | None, + result: SendResult | None, + retry_delay: int | None, + ) -> None: + """Record an export outcome under the run's row lock. + + An attempt whose token was replaced records nothing. An acceptance that + arrives after its lease expired still records the destination's result. + """ + ... + + def recovery_checked(self, run_id: str, *, error_code: str | None) -> None: ... + + def pending_ids(self, limit: int) -> list[str]: ... + + def retry(self, *, run_id: str, acknowledge_possible_duplicates: bool) -> None: + """Release a failed or uncertain run for another export. + + Raises: + TraceCollectionNotFoundError: Nothing was collected. + TraceCollectionConflictError: The run is accepted, not retryable, or needs the duplicate acknowledgement. + """ + ... + + def purge(self, *, retention_days: int, limit: int) -> int: ... diff --git a/api/services/tracing/delivery.py b/api/services/tracing/delivery.py index 592e9158ef6301..9e91e3e0a80d22 100644 --- a/api/services/tracing/delivery.py +++ b/api/services/tracing/delivery.py @@ -8,15 +8,11 @@ rotated or removed credential takes effect immediately. """ -import json import logging -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass -from datetime import UTC, datetime from typing import Any, Protocol -from uuid import uuid4 -from celery import current_app as celery_app from dify_trace_sdk.config import ProviderConfig from dify_trace_sdk.egress import Egress from dify_trace_sdk.errors import ExportOutcomeUnknownError, ExportRejectedError, ExportRetryableError @@ -25,57 +21,63 @@ from dify_trace_sdk.provider import TraceProvider from opentelemetry import metrics from pydantic import BaseModel, ConfigDict, ValidationError -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from configs import dify_config -from enterprise.telemetry.enterprise_trace import EnterpriseOtelTrace -from enterprise.telemetry.exporter import is_enterprise_telemetry_enabled -from extensions.ext_database import db -from extensions.ext_storage import storage -from models.model import App, TraceAppConfig -from services.tracing.assembly import TraceAssembler + from services.tracing.canonical import CanonicalTraceBuilder -from services.tracing.egress import SSRFEgress -from services.tracing.emitter import TRACE_QUEUE, payload_key -from services.tracing.provider_configs import ProviderConfigUnreadableError, decode_config +from services.tracing.errors import ( + ProviderConfigUnreadableError, + TraceCollectorUnavailableError, + TraceProviderUnavailableError, + UnknownTraceProviderError, +) +from services.tracing.providers import ProviderConfigCodec, TraceProviderCatalog from services.tracing.records import BaseTraceInfo, MessageTraceInfo, WorkflowTraceInfo -from services.tracing.registry import TraceProviderUnavailableError, UnknownTraceProviderError, provider_class -from services.tracing.service import application_tracing_service logger = logging.getLogger(__name__) -RETRY_TASK = "tasks.trace_delivery_task.retry_trace_export" - _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 RecordAssembler(Protocol): +class TraceRecordSource(Protocol): def assemble(self, event: TraceEvent) -> BaseTraceInfo: - """Raises SourceNotFoundError, or SourceNotReadyError while the source is not persisted yet.""" + """Read the event's record through its tenant and app. + + 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, info: BaseTraceInfo) -> bool: + 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, info: WorkflowTraceInfo) -> WorkflowTraceInfo: + 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.""" ... -class TraceCollectorUnavailableError(Exception): - """The Agent trace collector could not decide ownership; the event must be redelivered.""" - - -@dataclass(frozen=True) -class Destination: - provider: TraceProvider[Any] - config: ProviderConfig +type EnterpriseTraceSink = Callable[[BaseTraceInfo], None] class ExportJob(BaseModel): @@ -90,6 +92,18 @@ class ExportJob(BaseModel): 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)) @@ -98,18 +112,25 @@ class TraceDeliveryService: def __init__( self, *, - session_factory: sessionmaker[Session], - assembler: RecordAssembler, + records: TraceRecordSource, + destinations: TraceDestinationQuery, + providers: TraceProviderCatalog, + codec: ProviderConfigCodec, builder: CanonicalTraceBuilder, egress: Egress, - enterprise_sink: Callable[[BaseTraceInfo], None] | None, + retries: ExportRetryQueue, + enterprise_sink: EnterpriseTraceSink | None, collector: TraceCollector | None, max_export_attempts: int, ) -> None: - self._sessions = session_factory - self._assembler = assembler + """``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 @@ -118,7 +139,7 @@ def deliver(self, event: TraceEvent, *, enterprise_delivered: bool) -> None: """Raises SourceNotReadyError (redeliver later) or TraceCollectorUnavailableError.""" kind = event.payload.kind try: - record = self._assembler.assemble(event) + 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"}) @@ -130,11 +151,11 @@ def deliver(self, event: TraceEvent, *, enterprise_delivered: bool) -> None: 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(record): + 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(record) + 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: @@ -188,7 +209,9 @@ def export(self, job: ExportJob, destination: Destination) -> None: 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._schedule_retry(job.model_copy(update={"attempt": job.attempt + 1})) + 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 @@ -208,64 +231,16 @@ def export(self, job: ExportJob, destination: Destination) -> None: def destination(self, tenant_id: str, app_id: str) -> Destination | None: """The app's current provider destination, read through the app's workspace.""" - with self._sessions() as session: - app_tracing = session.scalar(select(App.tracing).where(App.id == app_id, App.tenant_id == tenant_id)) - provider_id = self._enabled_provider(app_tracing) - if provider_id is None: - return None - stored = session.scalar( - select(TraceAppConfig.tracing_config).where( - TraceAppConfig.app_id == app_id, - TraceAppConfig.tracing_provider == provider_id, - TraceAppConfig.is_active.is_(True), - ) - ) - if not isinstance(stored, dict): + stored = self._destinations.active_destination(tenant_id=tenant_id, app_id=app_id) + if stored is None: return None try: - provider = provider_class(provider_id) - config = decode_config(tenant_id, provider, stored) + 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", provider_id, app_id) + 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", provider_id, app_id) + 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) - - @staticmethod - def _enabled_provider(app_tracing: str | None) -> str | None: - if not app_tracing: - return None - try: - policy = json.loads(app_tracing) - except ValueError: - return None - if not isinstance(policy, dict) or policy.get("enabled") is not True: - return None - provider_id = policy.get("tracing_provider") - return provider_id if isinstance(provider_id, str) else None - - @staticmethod - def _schedule_retry(job: ExportJob) -> 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,), queue=TRACE_QUEUE, countdown=retry_delay_seconds(job.attempt - 1)) - - -def _enterprise_sink(record: BaseTraceInfo) -> None: - EnterpriseOtelTrace().trace(record) - - -def trace_delivery_service() -> TraceDeliveryService: - sessions = sessionmaker(bind=db.engine, expire_on_commit=False) - enterprise = is_enterprise_telemetry_enabled() - return TraceDeliveryService( - session_factory=sessions, - assembler=TraceAssembler(sessions, files_url=dify_config.FILES_URL, enterprise_details=enterprise), - builder=CanonicalTraceBuilder(), - egress=SSRFEgress(), - enterprise_sink=_enterprise_sink if enterprise else None, - collector=application_tracing_service() if dify_config.OPS_AGENT_TRACE_ENABLED else None, - max_export_attempts=dify_config.OPS_TRACE_EXPORT_MAX_ATTEMPTS, - ) diff --git a/api/services/tracing/emitter.py b/api/services/tracing/emitter.py deleted file mode 100644 index 75fb7074c5b116..00000000000000 --- a/api/services/tracing/emitter.py +++ /dev/null @@ -1,129 +0,0 @@ -"""The request-path side of the trace pipeline. - -Deciding whether to trace costs nothing extra: generators already hold the -app's tracing policy, so no configuration is decrypted and no provider is -loaded while a user waits. Events go straight to the delivery worker; large -ones travel through storage instead of the broker. -""" - -import json -import logging -from collections.abc import Callable, Mapping -from datetime import datetime -from typing import Any, override - -from celery import current_app as celery_app -from dify_trace_sdk.emitter import TraceEmitter, TraceScope -from dify_trace_sdk.events import TraceCorrelation, TraceEvent -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.helper.trace_id_helper import ParentTraceContext -from extensions.ext_storage import storage -from models.model import App - -logger = logging.getLogger(__name__) - -DELIVER_TASK = "tasks.trace_delivery_task.deliver_trace_event" -TRACE_QUEUE = "ops_trace" -PAYLOAD_PREFIX = "trace_payloads" -INLINE_EVENT_BYTES = 256 * 1024 -# Node executions may be saved by another worker after their run row; give them time to land. -WORKFLOW_SETTLE_SECONDS = 5 - - -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" - - -def tracing_policy_enabled(app_tracing: str | None) -> bool: - if not app_tracing: - return False - try: - policy = json.loads(app_tracing) - except ValueError: - return False - return ( - isinstance(policy, dict) and policy.get("enabled") is True and isinstance(policy.get("tracing_provider"), str) - ) - - -def correlation_from_extras(extras: Mapping[str, Any]) -> TraceCorrelation: - 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, - ) - - -class CeleryTraceEmitter(TraceEmitter): - @override - def emit(self, event: TraceEvent) -> None: - encoded = event.model_dump_json() - if len(encoded.encode()) > INLINE_EVENT_BYTES: - key = payload_key(event.occurred_at, "events", event.event_id) - storage.save(key, encoded.encode()) - envelope = {"payload_key": key} - else: - envelope = {"event": encoded} - countdown = WORKFLOW_SETTLE_SECONDS if event.payload.kind == "workflow_run" else 0 - celery_app.send_task(DELIVER_TASK, args=(envelope,), queue=TRACE_QUEUE, countdown=countdown) - - -class TraceScopes: - """Create trace scopes for executions whose events someone consumes.""" - - def __init__( - self, - *, - emitter: TraceEmitter, - session_factory: sessionmaker[Session], - enterprise_enabled: Callable[[], bool], - ) -> None: - self._emitter = emitter - self._sessions = session_factory - self._enterprise_enabled = enterprise_enabled - - def __call__( - self, - *, - tenant_id: str, - app_id: str | None, - app_tracing: str | None, - user_id: str | None, - extras: Mapping[str, Any], - ) -> TraceScope | None: - if not (tracing_policy_enabled(app_tracing) 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_from_extras(extras), - ) - - def for_app( - self, *, tenant_id: str, app_id: str | None, user_id: str | None, extras: Mapping[str, Any] - ) -> TraceScope | None: - """A scope for callers that do not hold the app row; tracing failures stay out of the request.""" - app_tracing = None - if app_id is not None: - try: - with self._sessions() as session: - app_tracing = session.scalar( - select(App.tracing).where(App.id == app_id, App.tenant_id == tenant_id) - ) - except Exception: - logger.warning("Tracing policy lookup failed: app_id=%s", app_id, exc_info=True) - return None - return self(tenant_id=tenant_id, app_id=app_id, app_tracing=app_tracing, user_id=user_id, extras=extras) 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/provider_configs.py b/api/services/tracing/provider_configs.py index d76285ee1f28d5..347c0031201dd0 100644 --- a/api/services/tracing/provider_configs.py +++ b/api/services/tracing/provider_configs.py @@ -1,4 +1,4 @@ -"""Store, verify, and present tracing provider configurations. +"""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 @@ -8,7 +8,8 @@ import logging from collections.abc import Mapping -from typing import Any, override +from dataclasses import dataclass +from typing import Any, Literal, override from dify_trace_sdk.config import ProviderConfig from dify_trace_sdk.egress import Egress @@ -16,7 +17,6 @@ from dify_trace_sdk.provider import TraceProvider from pydantic import ValidationError -from core.helper.encrypter import batch_decrypt_token, encrypt_token, obfuscated_token from services.app_tracing_config_service import ( AppTracingConfigInvalidConfigurationError, AppTracingConfigInvalidProviderError, @@ -26,78 +26,38 @@ PresentedTracingConfig, TracingConfigProviderGateway, ) -from services.tracing.registry import ( - SUPPORTED_PROVIDERS, +from services.tracing.errors import ( + ProviderConfigUnreadableError, TraceProviderUnavailableError, UnknownTraceProviderError, - provider_class, +) +from services.tracing.providers import ( + SUPPORTED_PROVIDERS, + ProviderConfigCodec, + TraceProviderCatalog, + configured_secrets, ) logger = logging.getLogger(__name__) -class ProviderConfigUnreadableError(Exception): - """A stored configuration cannot be decrypted or no longer validates.""" - - -def resolve_provider(provider_id: str) -> type[TraceProvider[Any]]: +def resolve_provider(providers: TraceProviderCatalog, provider_id: str) -> type[TraceProvider[Any]]: try: - return provider_class(provider_id) + return providers.resolve(provider_id) except UnknownTraceProviderError: raise AppTracingConfigInvalidProviderError(provider_id) from None except TraceProviderUnavailableError: raise AppTracingConfigProviderUnavailableError from None -def encode_config(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] = encrypt_token(tenant_id, value) - return stored - - -def decode_config(tenant_id: str, provider: type[TraceProvider[Any]], stored: Mapping[str, Any]) -> ProviderConfig: - """Decrypt and validate a stored configuration for verification or delivery.""" - 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 = batch_decrypt_token(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 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)) - - def _comparable(value: object) -> object: return value.strip().rstrip("/") if isinstance(value, str) else value -def _left_unchanged(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 == obfuscated_token(current_value) - ) - - class ProviderConfigGateway(TracingConfigProviderGateway): - def __init__(self, egress: Egress) -> None: + def __init__(self, *, providers: TraceProviderCatalog, codec: ProviderConfigCodec, egress: Egress) -> None: + self._providers = providers + self._codec = codec self._egress = egress @override @@ -110,9 +70,9 @@ def validate_provider(self, tracing_provider: str) -> None: def validate_saved_config( self, *, workspace_id: str, tracing_provider: str, tracing_config: dict[str, Any] ) -> None: - provider = resolve_provider(tracing_provider) + provider = resolve_provider(self._providers, tracing_provider) try: - decode_config(workspace_id, provider, tracing_config) + self._codec.decode(workspace_id, provider, tracing_config) except ProviderConfigUnreadableError: raise AppTracingConfigProcessingError from None @@ -120,7 +80,7 @@ def validate_saved_config( def prepare_new_config( self, *, workspace_id: str, tracing_provider: str, tracing_config: dict[str, Any] ) -> dict[str, Any]: - provider = resolve_provider(tracing_provider) + provider = resolve_provider(self._providers, tracing_provider) return self._verified(workspace_id, provider, self._submitted(provider, tracing_config, None)) @override @@ -132,11 +92,11 @@ def prepare_updated_config( tracing_config: dict[str, Any], current_tracing_config: dict[str, Any] | None, ) -> dict[str, Any]: - provider = resolve_provider(tracing_provider) + provider = resolve_provider(self._providers, tracing_provider) current = None if current_tracing_config is not None: try: - current = decode_config(workspace_id, provider, current_tracing_config) + 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 @@ -148,7 +108,7 @@ def present_config( ) -> PresentedTracingConfig: if tracing_config is None: raise AppTracingConfigProcessingError - provider = resolve_provider(tracing_provider) + provider = resolve_provider(self._providers, tracing_provider) model = provider.config_model hidden = model.secret_fields() | model.derived_fields() try: @@ -164,9 +124,8 @@ def present_config( values["project_url"] = provider(self._egress).project_url(config) return PresentedTracingConfig(values=values, configured_secrets=configured_secrets(provider, tracing_config)) - @staticmethod def _submitted( - provider: type[TraceProvider[Any]], submitted: Mapping[str, Any], current: ProviderConfig | None + self, provider: type[TraceProvider[Any]], submitted: Mapping[str, Any], current: ProviderConfig | None ) -> dict[str, Any]: """Merge a form submission over the current configuration. @@ -184,7 +143,7 @@ def _submitted( values: dict[str, Any] = {} for field in model.form_fields(): current_value = current_values.get(field.name) - if field.secret and _left_unchanged(submitted, field.name, current_value): + 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 @@ -192,6 +151,15 @@ def _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]: @@ -204,19 +172,48 @@ def _verified( except VerificationError as error: logger.warning("Tracing configuration verification failed: provider=%s code=%s", provider.id, error.code) raise AppTracingConfigVerificationFailedError from None - return encode_config(workspace_id, verified) - - -def capabilities_payload(provider_id: str) -> dict[str, bool | str]: - """What a provider's mapping supports, for the console. Describes the adapter, not delivery health.""" - capabilities = resolve_provider(provider_id).capabilities - return { - "provider": provider_id, - "agent_spans": True, - "session": True, - "native_usage": True, - "native_cost": capabilities.native_cost, - "trace_url": capabilities.trace_url, - "span_url": capabilities.span_url, - "export_ack": "accepted", - } + return self._codec.encode(workspace_id, verified) + + +@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.""" + + +@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.""" + + +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, + ) + + +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/providers.py b/api/services/tracing/providers.py new file mode 100644 index 00000000000000..b45ef871aa0fbe --- /dev/null +++ b/api/services/tracing/providers.py @@ -0,0 +1,104 @@ +"""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 typing import Any, 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)) diff --git a/api/services/tracing/records.py b/api/services/tracing/records.py index 72271371995598..0e1a911a916304 100644 --- a/api/services/tracing/records.py +++ b/api/services/tracing/records.py @@ -1,26 +1,59 @@ """Assembled trace records: one Dify operation with the data its sinks need. -Records live only inside the delivery worker. They may reference ORM rows -because they never cross a process or storage boundary; the canonical trace -built from them is what gets persisted for retries. +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 +from pydantic import BaseModel, ConfigDict, StrictStr, ValidationError -from core.helper.trace_id_helper import ParentTraceContext from graphon.entities import WorkflowNodeExecution -from models.model import Message -from models.workflow import WorkflowRun + + +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: Message | 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 @@ -29,7 +62,7 @@ class BaseTraceInfo(BaseModel): trace_id: str | None = None """The caller-supplied trace ID, if any.""" - model_config = ConfigDict(arbitrary_types_allowed=True, protected_namespaces=(), frozen=True) + model_config = ConfigDict(protected_namespaces=(), frozen=True) @property def resolved_trace_id(self) -> str | None: @@ -46,21 +79,21 @@ def resolved_trace_id(self) -> str | None: @property def resolved_parent_context(self) -> tuple[str | None, str | None]: """(outer workflow_run_id, outer node_execution_id) of a nested workflow.""" - match self.metadata.get("parent_trace_context"): - case ParentTraceContext() as context: - pass - case Mapping() as raw: - try: - context = ParentTraceContext.model_validate(raw) - except ValueError: - return None, None - case _: - return None, None + 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): - workflow_data: WorkflowRun + start_time: datetime + """When the run was created.""" + end_time: datetime + """When the run finished; only finished runs are traced.""" node_executions: Sequence[WorkflowNodeExecution] conversation_id: str | None message_started_at: datetime | None @@ -84,7 +117,7 @@ class WorkflowTraceInfo(BaseTraceInfo): class MessageTraceInfo(BaseTraceInfo): - message_data: Message + message_data: TraceMessage conversation_model: str message_tokens: int answer_tokens: int diff --git a/api/services/tracing/recovery.py b/api/services/tracing/recovery.py index 26943a3c4659bb..de02588f35c247 100644 --- a/api/services/tracing/recovery.py +++ b/api/services/tracing/recovery.py @@ -5,27 +5,23 @@ Message is complete. Recovery never invents answers, usage, or success states. """ -import json import logging -import math from datetime import UTC, datetime, timedelta -from typing import Protocol from uuid import uuid4 from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent, TracePayload, WorkflowRunCompleted -from sqlalchemy import select -from sqlalchemy.exc import SQLAlchemyError -from configs import dify_config from graphon.enums import WorkflowExecutionStatus -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 -from models.workflow import WorkflowRun -from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError -from services.tracing.assembly import SourceNotFoundError, SourceNotReadyError, TraceAssembler -from services.tracing.records import BaseTraceInfo, MessageTraceInfo, WorkflowTraceInfo +from services.tracing.collection import ApplicationTraceStore, TraceRun, TraceRunState +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 @@ -37,47 +33,7 @@ } -class RecoveryTarget(Protocol): - repository: ApplicationTraceRepository - - @property - def assembler(self) -> TraceAssembler: ... - - 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 _event(run: ApplicationTraceRun, payload: TracePayload, user_id: str | None) -> TraceEvent: +def _event(run: TraceRun, payload: TracePayload, user_id: str | None) -> TraceEvent: return TraceEvent( event_id=uuid4().hex, tenant_id=run.tenant_id, @@ -94,119 +50,102 @@ def _event(run: ApplicationTraceRun, payload: TracePayload, user_id: str | None) ) -def _message_task(repository: ApplicationTraceRepository, run: ApplicationTraceRun, now: datetime) -> TraceEvent | 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 _event(run, MessageCompleted(message_id=message.id, conversation_id=message.conversation_id), None) +def _message_event( + traces: ApplicationTraceStore, run: TraceRun, now: datetime, collection_timeout: int +) -> TraceEvent | None: + progress = traces.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") + pending_agent_stop = progress.agent_app and progress.agent_stop_requested + if progress.paused and not pending_agent_stop: + return None + invocations = traces.invocations(run.id) + 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 TraceCollectionError("agent_trace_collection_incomplete") + return None + 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 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 + if not (progress.explicit_error or progress.answer_saved): + if finalized and max(finalized) + timedelta(seconds=collection_timeout) < now: + raise TraceCollectionError("business_terminal_not_persisted") + 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 -) -> TraceEvent | None: - workflow = repository._source(tenant_id=run.tenant_id, app_id=run.app_id, kind="workflow", source_id=run.source_id) +def _workflow_event(traces: ApplicationTraceStore, run: TraceRun, now: datetime) -> TraceEvent | None: + workflow = traces.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 workflow.status not in _WORKFLOW_TERMINAL or workflow.finished_at is None: 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 _event( - run, - WorkflowRunCompleted(workflow_run_id=workflow.id, conversation_id=conversation_id, total_tokens=None), - workflow.created_by, - ) + conversation_id = traces.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: TraceRun, + *, + traces: ApplicationTraceStore, + records: TraceRecordSource, + collector: TraceCollector, + collection_timeout: int, + now: datetime, +) -> bool: + """Attempt one bounded recovery and rotate unresolved work behind other runs. -def recover_terminal(target: RecoveryTarget, run: ApplicationTraceRun) -> bool: - """Attempt one bounded recovery and rotate unresolved work behind other runs.""" + ``now`` is naive UTC, like the ledger's timestamps. + """ if ( run.outer_trace is not None - or run.state not in {"collecting", "ready"} + or run.state not in {TraceRunState.COLLECTING, TraceRunState.READY} or run.source_kind not in {"message", "workflow"} ): return False error_code: str | None = None try: - event = (_message_task if run.source_kind == "message" else _workflow_task)( - target.repository, run, naive_utc_now() + event = ( + _message_event(traces, run, now, collection_timeout) + if run.source_kind == "message" + else _workflow_event(traces, run, now) ) if event is None: return False try: - info = target.assembler.assemble(event) + info = records.assemble(event) except SourceNotReadyError: return False 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: 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 @@ -215,4 +154,4 @@ 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) + traces.recovery_checked(run.id, error_code=error_code) diff --git a/api/services/tracing/registry.py b/api/services/tracing/registry.py index 988b11211fdabe..02f847ae13941e 100644 --- a/api/services/tracing/registry.py +++ b/api/services/tracing/registry.py @@ -1,41 +1,19 @@ """Discover installed trace providers through ``dify.trace_providers`` entry points. -Dify supports a fixed set of provider IDs; a distribution may install any -subset of their packages. An unknown ID is a client error, a known but -uninstalled one is a deployment fact. Neither ever raises into a business -request: callers translate both. +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 +from typing import Any, override from dify_trace_sdk.provider import TraceProvider -ENTRY_POINT_GROUP = "dify.trace_providers" - -SUPPORTED_PROVIDERS = ( - "langfuse", - "langsmith", - "opik", - "weave", - "arize", - "phoenix", - "aliyun", - "mlflow", - "databricks", - "tencent", -) - - -class UnknownTraceProviderError(Exception): - def __init__(self, provider_id: str) -> None: - super().__init__(f"Unknown tracing provider: {provider_id}") +from services.tracing.errors import TraceProviderUnavailableError, UnknownTraceProviderError +from services.tracing.providers import SUPPORTED_PROVIDERS, TraceProviderCatalog - -class TraceProviderUnavailableError(Exception): - def __init__(self, provider_id: str) -> None: - super().__init__(f"Tracing provider {provider_id} is not installed in this deployment") +ENTRY_POINT_GROUP = "dify.trace_providers" @functools.cache @@ -44,9 +22,7 @@ def _entry_points() -> dict[str, EntryPoint]: @functools.cache -def provider_class(provider_id: str) -> type[TraceProvider[Any]]: - if provider_id not in SUPPORTED_PROVIDERS: - raise UnknownTraceProviderError(provider_id) +def _load(provider_id: str) -> type[TraceProvider[Any]]: entry_point = _entry_points().get(provider_id) if entry_point is None: raise TraceProviderUnavailableError(provider_id) @@ -59,13 +35,11 @@ def provider_class(provider_id: str) -> type[TraceProvider[Any]]: return loaded -def installed_providers() -> tuple[str, ...]: - """Supported providers whose packages load in this deployment, in display order.""" - installed: list[str] = [] - for provider_id in SUPPORTED_PROVIDERS: - try: - provider_class(provider_id) - except TraceProviderUnavailableError: - continue - installed.append(provider_id) - return tuple(installed) +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 fec0ac560caf83..f8adfe2e4ce496 100644 --- a/api/services/tracing/runtime_bridge.py +++ b/api/services/tracing/runtime_bridge.py @@ -1,168 +1,133 @@ """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.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) - -class RuntimeApplicationTracingHooks: - """Stateless app-scoped adapter; database services are created only on invocation.""" - - __slots__ = () + @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 - prepare_agent_app = staticmethod(prepare_agent_app) - prepare_workflow_agent = staticmethod(prepare_workflow_agent) - freeze_workflow_run = staticmethod(freeze_workflow_run) - 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/service.py b/api/services/tracing/service.py index c1aec1f470d5cc..ced9a3d47e6689 100644 --- a/api/services/tracing/service.py +++ b/api/services/tracing/service.py @@ -12,7 +12,7 @@ from collections.abc import Generator from contextlib import contextmanager from datetime import UTC, datetime, timedelta -from typing import Any, Literal +from typing import Any, Literal, override from uuid import uuid4 from dify_agent.layers.execution_context import ApplicationTracingContext @@ -24,21 +24,33 @@ from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace from dify_trace_sdk.provider import TraceProvider from pydantic import ValidationError -from sqlalchemy import select -from sqlalchemy.orm import sessionmaker - -from configs import dify_config -from extensions.ext_database import db -from models import App, AppMode, Message, Workflow, WorkflowNodeExecutionModel, WorkflowRun -from models.application_trace import ApplicationTraceRun -from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError -from services.app_tracing_config_service import ( - AppTracingConfigInvalidProviderError, - AppTracingConfigProviderUnavailableError, -) -from services.tracing.assembly import TraceAssembler + +from libs.datetime_utils import naive_utc_now from services.tracing.canonical import CanonicalTraceBuilder -from services.tracing.egress import SSRFEgress +from services.tracing.collection import ( + CLOSED_STATES, + AgentTraceSettings, + ApplicationTraceStore, + ScheduleTraceExport, + SourceMessage, + SourceNodeExecution, + TraceApp, + TraceAppMode, + TraceRun, + TraceRunState, + TraceRunStatus, + TraceSource, + TraceSourceKind, + workflow_has_agent_v2, +) +from services.tracing.delivery import TraceCollector, TraceRecordSource +from services.tracing.errors import ( + ProviderConfigUnreadableError, + TraceCollectionError, + TraceCollectionNotFoundError, + TraceProviderUnavailableError, + UnknownTraceProviderError, +) from services.tracing.normalization import ( legacy_outer_trace, normalize_agent_span, @@ -46,14 +58,15 @@ remap_outer_trace, sanitize, ) -from services.tracing.provider_configs import ProviderConfigUnreadableError, capabilities_payload, decode_config -from services.tracing.records import BaseTraceInfo, MessageTraceInfo, WorkflowTraceInfo -from services.tracing.recovery import agent_stop_requested, recover_terminal -from services.tracing.registry import TraceProviderUnavailableError, UnknownTraceProviderError, provider_class -from services.tracing.runtime_bridge import workflow_has_agent_v2 +from services.tracing.provider_configs import TraceProviderSupport, support_of +from services.tracing.providers import ProviderConfigCodec, TraceProviderCatalog +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): @@ -66,33 +79,42 @@ def _message_inputs_are_empty(inputs: object) -> bool: return inputs is None or isinstance(inputs, (str, dict, list)) and not inputs -class ApplicationTracingService: +class ApplicationTracingService(TraceCollector): def __init__( self, - repository: ApplicationTraceRepository, *, - assembler: TraceAssembler, + traces: ApplicationTraceStore, + records: TraceRecordSource, builder: CanonicalTraceBuilder, + providers: TraceProviderCatalog, + codec: ProviderConfigCodec, egress: Egress, + schedule_export: ScheduleTraceExport, + settings: AgentTraceSettings, ) -> None: - self.repository = repository - self.assembler = assembler + self._traces = traces + 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; ``None`` when the source is not collected.""" + if not self._settings.enabled: return None - run = self.repository.start( + run = self._traces.start( tenant_id=tenant_id, app_id=app_id, source_kind=source_kind, @@ -100,7 +122,7 @@ 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 prepare_invocation( self, @@ -108,13 +130,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, @@ -126,7 +148,7 @@ def prepare_invocation( ) if run is None or not run.capture_agent: return None - invocation = self.repository.invocation( + invocation = self._traces.invocation( run_id=run.id, agent_id=agent_id, node_execution_id=node_execution_id, attempt=attempt, call_key=call_key ) return ApplicationTracingContext( @@ -139,13 +161,11 @@ 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, + accepted, run_id = self._traces.ingest( + batch, max_spans=self._settings.max_spans, max_bytes=self._settings.max_bytes ) # Persistence commits before scheduling. Beat reconciles a broker outage. - self.schedule(run_id) + self._schedule(run_id) return accepted def capture_agent_app_terminal( @@ -163,15 +183,13 @@ 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) - run = self.repository.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": + source = self._traces.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._traces.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 is TraceRunState.NOT_COLLECTED: return - info = self.assembler.assemble( + info = self._records.assemble( TraceEvent( event_id=uuid4().hex, tenant_id=tenant_id, @@ -194,54 +212,48 @@ def capture_agent_app_terminal( # Early errors may precede formatted-message persistence. Preserve only # the known user query on the root; this is not an inferred LLM prompt. inputs = {"query": info.message_data.query} - if output is not None: - # This is a detached trace snapshot, not a write to the business row. - info.message_data.answer = output + message = info.message_data if output is None else info.message_data.model_copy(update={"answer": output}) self.capture_terminal( - info.model_copy( + 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}", } - ) + ), ) - @staticmethod - def schedule(run_id: str) -> None: - from tasks.application_trace_task import export_application_trace + def reconcile(self) -> None: + """Schedule unfinished runs again, behind a broker outage, and purge runs past their retention.""" + for run_id in self._traces.pending_ids(_RECONCILE_BATCH): + self._schedule(run_id) + self._traces.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 _run_for_source( - self, *, tenant_id: str, app_id: str, source_kind: str, source_id: str - ) -> ApplicationTraceRun | None: - run = self.repository.source_run( - tenant_id=tenant_id, app_id=app_id, source_kind=source_kind, source_id=source_id - ) + self, *, tenant_id: str, app_id: str, source_kind: TraceSourceKind, source_id: str + ) -> TraceRun | None: + run = self._traces.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._traces.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._traces.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: + 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( @@ -249,64 +261,40 @@ def _managed_owner(self, app: App, message: Message | None = None, workflow_run_ ): 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._traces.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._traces.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) + graph = ( + self._traces.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 _app(self, app_id: str) -> App | None: - with self.repository.sessions() as session: - return session.scalar(select(App).where(App.id == app_id)) - - def link_managed_parent(self, info: WorkflowTraceInfo) -> WorkflowTraceInfo: + @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.""" - tenant_id = info.tenant_id parent_run_id, parent_node_id = info.resolved_parent_context if not parent_run_id: return info - 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, - ) - ) + parent = self._traces.workflow_run_trace(tenant_id=tenant_id, workflow_run_id=parent_run_id) if parent is None: return info - if ( - self.repository._source(tenant_id=tenant_id, app_id=parent.app_id, kind="workflow", source_id=parent_run_id) - is None - ): + if self._traces.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) @@ -317,8 +305,10 @@ def link_managed_parent(self, info: WorkflowTraceInfo) -> WorkflowTraceInfo: metadata["linked_parent_trace_id"] = parent.trace_id return info.model_copy(update={"metadata": metadata}) - def capture_terminal(self, info: BaseTraceInfo) -> bool: + @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" @@ -328,19 +318,22 @@ def capture_terminal(self, info: BaseTraceInfo) -> bool: else: # 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 = info.message_data - owner_app_id = message.app_id if message is not None else info.metadata.get("app_id") - app = self._app(owner_app_id) if isinstance(owner_app_id, str) and owner_app_id else None - return app is not None and self._managed_owner(app, message) - 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" + 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._traces.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._traces.app(tenant_id=tenant_id, app_id=app_id) + if app is None: + return False + message = ( + self._traces.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": @@ -349,7 +342,7 @@ def capture_terminal(self, info: BaseTraceInfo) -> bool: if run is None: if not managed: return False - run = self.repository.record_initialization_failure( + run = self._traces.record_initialization_failure( tenant_id=app.tenant_id, app_id=app.id, source_kind=source_kind, @@ -361,13 +354,13 @@ def capture_terminal(self, info: BaseTraceInfo) -> bool: ) if run is None: return True - if run.state in {"not_collected", "failed", "uncertain", "accepted", "exporting"}: + if run.state is TraceRunState.NOT_COLLECTED or run.state in CLOSED_STATES: 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 agent_stop_requested(message) + and message.agent_stop_requested ): # An already queued consumer notification may arrive after a # stop. Its ordinary Message snapshot cannot own the runtime @@ -381,15 +374,15 @@ def capture_terminal(self, info: BaseTraceInfo) -> bool: trace, trace_id=run.trace_id, capture_content=run.capture_content, synthetic_span_ids=synthetic ) 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._traces.terminal(run.id, trace.model_dump(mode="json"), trace.root.span_id) + self._schedule(run.id) return True - def settle_preview(self, run: ApplicationTraceRun) -> None: + def _settle_preview(self, run: TraceRun) -> 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"}: + node = self._traces.node_execution(tenant_id=run.tenant_id, app_id=run.app_id, node_execution_id=run.source_id) + if node is None or node.status in {"running", "pending", "paused"}: return error = node.error started_at = node.created_at.replace(tzinfo=UTC) @@ -415,9 +408,9 @@ def settle_preview(self, run: ApplicationTraceRun) -> None: parent=None, spans=(span,), ) - self.repository.terminal(run.id, trace.model_dump(mode="json"), run.root_span_id) + self._traces.terminal(run.id, trace.model_dump(mode="json"), run.root_span_id) - def assemble(self, run: ApplicationTraceRun) -> Trace: + def _assemble(self, run: TraceRun) -> Trace: if run.outer_trace is None: raise TraceCollectionError("tracing_terminal_missing") # Runs captured before the upgrade stored the previous trace shape, named by its root span. @@ -426,7 +419,7 @@ def assemble(self, run: ApplicationTraceRun) -> Trace: if "root_span_id" in run.outer_trace else Trace.model_validate(run.outer_trace) ) - invocations, batches = self.repository.collected(run.id) + invocations, batches = self._traces.collected(run.id) invocation_by_id = {inv.id: inv for inv in invocations} spans = list(outer.spans) internal_by_invocation: dict[str, dict[str, Span]] = {inv.id: {} for inv in invocations} @@ -500,12 +493,12 @@ def _export_lease_heartbeat(self, run_id: str, token: str) -> Generator[None]: after lease expiry unless a manual retry has already replaced its token. """ 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._traces.renew_export_lease(run_id, token): return except Exception as error: logger.warning( @@ -533,18 +526,25 @@ 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) + record = self._traces.get(run_id) if record is None: return - recover_terminal(self, record) - record = self.repository.get(run_id) + recover_terminal( + record, + traces=self._traces, + records=self._records, + collector=self, + collection_timeout=self._settings.collection_timeout_seconds, + now=naive_utc_now(), + ) + record = self._traces.get(run_id) if record is None: return - self.settle_preview(record) - claim = self.repository.claim( + self._settle_preview(record) + claim = self._traces.claim( run_id, - collection_timeout=dify_config.OPS_AGENT_TRACE_COLLECTION_TIMEOUT, - lease_seconds=dify_config.OPS_AGENT_TRACE_EXPORT_LEASE_SECONDS, + collection_timeout=self._settings.collection_timeout_seconds, + lease_seconds=self._settings.export_lease_seconds, ) if claim is None: return @@ -552,32 +552,33 @@ def deliver(self, run_id: str) -> None: with self._export_lease_heartbeat(run_id, token): try: try: - trace = self.assemble(run) + trace = self._assemble(run) 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 provider = self._frozen_provider(run) - result = provider.send(decode_config(run.tenant_id, type(provider), run.provider_config), trace) + result = provider.send(self._codec.decode(run.tenant_id, type(provider), run.provider_config), trace) except ExportOutcomeUnknownError: - self.repository.complete(run_id, token, state="uncertain", error_code="provider_acceptance_unknown") + self._complete_failed(run_id, token, TraceRunState.UNCERTAIN, "provider_acceptance_unknown") except ExportRetryableError: - if run.attempts >= dify_config.OPS_AGENT_TRACE_MAX_ATTEMPTS: - self.repository.complete(run_id, token, state="failed", error_code="provider_retry_exhausted") + if run.attempts >= self._settings.max_export_attempts: + self._complete_failed(run_id, token, TraceRunState.FAILED, "provider_retry_exhausted") else: - self.repository.complete( + self._traces.complete( run_id, token, - state="retrying", + state=TraceRunState.RETRYING, error_code="provider_temporarily_unavailable", + result=None, 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) + self._complete_failed(run_id, token, TraceRunState.FAILED, error.code) except ProviderConfigUnreadableError: - self.repository.complete(run_id, token, state="failed", error_code="tracing_config_unreadable") + self._complete_failed(run_id, token, TraceRunState.FAILED, "tracing_config_unreadable") except ExportRejectedError: - self.repository.complete(run_id, token, state="failed", error_code="provider_export_rejected") + self._complete_failed(run_id, token, TraceRunState.FAILED, "provider_export_rejected") except Exception as error: logger.warning( "Application trace export failed: run_id=%s provider=%s error_type=%s", @@ -585,34 +586,44 @@ 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._complete_failed(run_id, token, TraceRunState.FAILED, "provider_export_rejected") else: - self.repository.complete(run_id, token, state="accepted", result=result) + self._traces.complete( + run_id, token, state=TraceRunState.ACCEPTED, error_code=None, result=result, retry_delay=None + ) - def _frozen_provider(self, run: ApplicationTraceRun) -> TraceProvider[Any]: + def _complete_failed(self, run_id: str, token: str, state: TraceRunState, error_code: str) -> None: + self._traces.complete(run_id, token, state=state, error_code=error_code, result=None, retry_delay=None) + + def _frozen_provider(self, run: TraceRun) -> TraceProvider[Any]: try: - return provider_class(run.provider)(self._egress) + return self._providers.resolve(run.provider)(self._egress) except (UnknownTraceProviderError, TraceProviderUnavailableError): raise TraceCollectionError("tracing_provider_unavailable") from None - @staticmethod - def _frozen_capabilities(provider_id: str) -> dict[str, bool | str] | None: + def _frozen_capabilities(self, provider_id: str) -> TraceProviderSupport | None: try: - return capabilities_payload(provider_id) - except (AppTracingConfigInvalidProviderError, AppTracingConfigProviderUnavailableError): + return support_of(self._providers.resolve(provider_id)) + except (UnknownTraceProviderError, TraceProviderUnavailableError): return None - def _frozen_project_url(self, run: ApplicationTraceRun) -> str | None: + def _frozen_project_url(self, run: TraceRun) -> str | None: """Present the run's own destination, never the app's current one.""" try: - provider = provider_class(run.provider) + 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, @@ -622,75 +633,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._traces.message(tenant_id=tenant_id, app_id=app_id, message_id=source_id) + case "workflow": + return self._traces.workflow_run(tenant_id=tenant_id, app_id=app_id, workflow_run_id=source_id) + case "node": + return self._traces.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._traces.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._traces.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._traces.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": 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, - "queryable": None, - "capabilities": self._frozen_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, @@ -699,18 +737,8 @@ 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: - sessions = sessionmaker(bind=db.engine, expire_on_commit=False) - return ApplicationTracingService( - ApplicationTraceRepository(sessions), - assembler=TraceAssembler(sessions, files_url=dify_config.FILES_URL, enterprise_details=False), - builder=CanonicalTraceBuilder(), - egress=SSRFEgress(), - ) + self._traces.retry(run_id=run.id, acknowledge_possible_duplicates=acknowledge_possible_duplicates) + self._schedule(run.id) + return self._status_for_run(self._traces.get(run.id), source_kind, source_id, node_execution_id) diff --git a/api/tasks/application_trace_task.py b/api/tasks/application_trace_task.py index a7d68d69b8dd1a..6e1f6ae1b96007 100644 --- a/api/tasks/application_trace_task.py +++ b/api/tasks/application_trace_task.py @@ -2,21 +2,14 @@ from celery import shared_task -from configs import dify_config +from extensions.ext_application_services import application_services @shared_task(queue="ops_trace") def export_application_trace(run_id: str) -> None: - from services.tracing.service import application_tracing_service - - application_tracing_service().deliver(run_id) + application_services().tracing.agent_traces.deliver(run_id) @shared_task(queue="ops_trace") 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/trace_delivery_task.py b/api/tasks/trace_delivery_task.py index e8818a1b212777..b45ac0029642f4 100644 --- a/api/tasks/trace_delivery_task.py +++ b/api/tasks/trace_delivery_task.py @@ -13,10 +13,11 @@ 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.assembly import SourceNotFoundError, SourceNotReadyError -from services.tracing.delivery import ExportJob, TraceCollectorUnavailableError, trace_delivery_service -from services.tracing.emitter import PAYLOAD_PREFIX +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__) @@ -60,7 +61,9 @@ def deliver_trace_event(self, envelope: dict[str, Any]) -> None: _delete_quietly(payload_key) return try: - trace_delivery_service().deliver(event, enterprise_delivered=envelope.get("enterprise_delivered") is True) + 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: @@ -90,7 +93,7 @@ def retry_trace_export(job_key: str) -> None: _delete_quietly(job_key) return try: - trace_delivery_service().retry(job) + application_services().tracing.delivery.retry(job) finally: _delete_quietly(job_key) 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 62c9d6b3ba9f2c..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 @@ -46,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" @@ -278,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( @@ -300,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/controllers/console/app/test_message_suggested_questions.py b/api/tests/e2e_tests/contracts/controllers/console/app/test_message_suggested_questions.py index a56c4414d7988a..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 @@ -57,10 +57,10 @@ MessageSuggestedQuestionsService, ) from services.recommended_app_package_service import RecommendedAppPackageService -from services.tracing.provider_configs import ProviderConfigGateway 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 @@ -188,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=ProviderConfigGateway(RecordingEgress([], grpc_target=None)), + provider=deployment_provider_configs(RecordingEgress([], grpc_target=None)), ), oauth=_build_oauth_server_service(database_client=sqlite_session_factory, redis=redis), recommended_packages=RecommendedAppPackageService( @@ -205,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/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/unit_tests/conftest.py b/api/tests/unit_tests/conftest.py index 0158c381f3e137..ba969cda9d258b 100644 --- a/api/tests/unit_tests/conftest.py +++ b/api/tests/unit_tests/conftest.py @@ -25,9 +25,9 @@ from extensions import ext_redis from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.base import TypeBase -from services.tracing.provider_configs import ProviderConfigGateway 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 @@ -371,7 +371,7 @@ def app_services(sqlite_session_factory: sessionmaker[Session], rbac_domain: RBA tracing_configs=AppTracingConfigService( configs=SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory), # Unit tests never reach a tracing vendor; an unexpected request fails the test. - provider=ProviderConfigGateway(RecordingEgress([], grpc_target=None)), + 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 ee9adf85ebd052..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,6 +1,6 @@ import inspect from collections.abc import Callable -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from datetime import UTC, datetime from uuid import UUID @@ -19,6 +19,8 @@ from models.account import Account, AccountStatus, TenantAccountRole 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 """Agent tracing Console contracts and authorization stay separate from PlainApp.""" @@ -36,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 @@ -93,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 @@ -117,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 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 785600efd13fba..a5d58799475dba 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,17 +2,15 @@ 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 import pytest -from dify_trace_sdk.testing import RecordingEgress 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 @@ -23,11 +21,9 @@ 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.assembly import TraceAssembler -from services.tracing.canonical import CanonicalTraceBuilder -from services.tracing.service import ApplicationTracingService +from services.tracing.collection import TraceRunState, TraceRunStatus, TraceSourceKind +from services.tracing.errors import TraceCollectionConflictError, TraceCollectionError, TraceCollectionNotFoundError from tests.helpers.unit.config_override import apply_config_overrides CONTEXT = RequestContext("request", None, "account", "tenant") @@ -35,35 +31,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), - assembler=TraceAssembler(sqlite_session_factory, files_url="https://files.test", enterprise_details=False), - builder=CanonicalTraceBuilder(), - egress=RecordingEgress([], grpc_target=None), - ) - 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)) @@ -71,21 +77,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 @@ -100,18 +106,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)( @@ -132,12 +135,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 @@ -212,18 +214,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) @@ -232,6 +236,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 a43e8b16c434d4..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 @@ -50,9 +50,9 @@ AppTracingConfigService, AppTracingConfigVerificationFailedError, ) -from services.tracing.provider_configs import ProviderConfigGateway 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" @@ -156,7 +156,7 @@ def stored_configs(monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sess egress = RecordingEgress([], grpc_target=None) service = AppTracingConfigService( configs=SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory), - provider=ProviderConfigGateway(egress), + provider=deployment_provider_configs(egress), ) monkeypatch.setattr( ops_trace_module, "application_services", lambda: _ApplicationServices(app_tracing_configs=service) 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 9c998fad151a43..65100cd5cfbb2c 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,111 +1,178 @@ """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.application_trace_repository import ApplicationTraceRepository -from services.tracing.assembly import TraceAssembler +from repositories.trace_record_repository import TraceRecordRepository from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import AgentTraceSettings, 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"} -def _service(sessions: sessionmaker[Session], repository: ApplicationTraceRepository) -> ApplicationTracingService: - return ApplicationTracingService( - repository, - assembler=TraceAssembler(sessions, files_url="https://files.test", enterprise_details=False), + +@pytest.fixture +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"): + 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( + traces=collection.repository, + records=TraceRecordRepository(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) ) -@pytest.fixture -def client() -> Iterator[FlaskClient]: - 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() +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 = _service(sqlite_session_factory, 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 = _service(sqlite_session_factory, 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_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(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_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/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 7d734b9af7b6c3..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 @@ -19,7 +19,6 @@ from dify_trace_sdk.events import TraceEvent from flask import Flask from pydantic_ai.messages import BinaryContent, ImageUrl -from sqlalchemy.orm import sessionmaker from clients.agent_backend import ( DIFY_CONFIG_LAYER_ID, @@ -43,10 +42,9 @@ 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.emitter import TraceScopes -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) @@ -64,11 +62,6 @@ def emit(self, event: TraceEvent) -> None: self.events.append(event) -@pytest.fixture(autouse=True) -def _application_tracing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - - @pytest.fixture(autouse=True) def _no_runtime_agent_skills(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr( @@ -144,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 @@ -157,20 +150,13 @@ 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 = RecordingTracingHooks(tracing) app = Flask(__name__) bind_tracing( app, - scopes=TraceScopes(emitter=_TraceRecorder(), session_factory=sessionmaker(), enterprise_enabled=lambda: False), - hooks=RuntimeApplicationTracingHooks(), + scopes=kernel_trace_scopes(_TraceRecorder(), policies=NoTracingPolicies(), enterprise_enabled=False), + hooks=hooks, ) with app.app_context(): @@ -181,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/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 79fea3781db32d..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 @@ -11,7 +11,6 @@ from dify_agent.protocol import DIFY_AGENT_HISTORY_LAYER_ID, DIFY_AGENT_MODEL_LAYER_ID, DIFY_AGENT_OUTPUT_LAYER_ID from dify_agent.protocol.snapshot import SessionSnapshot from flask import Flask -from sqlalchemy.orm import sessionmaker from clients.agent_backend import ( DIFY_CONFIG_LAYER_ID, @@ -51,15 +50,10 @@ DeclaredOutputType, WorkflowNodeJobConfig, ) -from services.tracing.emitter import CeleryTraceEmitter, TraceScopes -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) @@ -248,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 @@ -263,22 +255,12 @@ 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_tracing( app, - scopes=TraceScopes( - emitter=CeleryTraceEmitter(), session_factory=sessionmaker(), enterprise_enabled=lambda: False - ), + scopes=kernel_trace_scopes(CeleryTraceEmitter(), policies=NoTracingPolicies(), enterprise_enabled=False), hooks=hooks, ) @@ -286,7 +268,7 @@ def prepare(value: WorkflowAgentRuntimeBuildContext, conversation_id: str | None 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/enterprise/telemetry/test_draft_trace.py b/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py index a62cf0e3eb8dd9..90502db2cca719 100644 --- a/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py +++ b/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py @@ -24,7 +24,7 @@ WorkflowNodeExecutionStatus, WorkflowNodeExecutionTriggeredFrom, ) -from services.tracing.assembly import TraceAssembler +from repositories.trace_record_repository import TraceRecordRepository from services.tracing.records import DraftNodeExecutionTrace # --------------------------------------------------------------------------- @@ -315,7 +315,7 @@ def test_none_outputs_uses_execution_outputs(self, mock_emit: MagicMock) -> None # --------------------------------------------------------------------------- # End-to-end token/model data flow: _build_node_execution_data → -# DraftNodeExecuted trace event → TraceAssembler → DraftNodeExecutionTrace +# DraftNodeExecuted trace event → TraceRecordRepository → DraftNodeExecutionTrace # --------------------------------------------------------------------------- @@ -526,7 +526,7 @@ def test_node_data_feeds_into_draft_node_execution_trace( (event,) = pipeline.events assert (event.tenant_id, event.app_id, event.user_id) == ("tenant-flow", "app-flow", "user-1") - trace_info = TraceAssembler( + trace_info = TraceRecordRepository( sqlite_session_factory, files_url="https://files.example.com", enterprise_details=False ).assemble(event) 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 763908953eab05..5aa4c413cefc6e 100644 --- a/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py +++ b/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py @@ -17,7 +17,6 @@ EnterpriseTelemetryHistogram, EnterpriseTelemetrySpan, ) -from models.workflow import WorkflowRun from services.tracing.records import ( DatasetRetrievalTraceInfo, DraftNodeExecutionTrace, @@ -27,10 +26,10 @@ PromptGenerationTraceInfo, SuggestedQuestionTraceInfo, ToolTraceInfo, + TraceMessage, WorkflowNodeTraceInfo, WorkflowTraceInfo, ) -from tests.helpers.unit.model_factories import make_message # --------------------------------------------------------------------------- # Fixtures @@ -61,15 +60,25 @@ def trace_handler(mock_exporter): _T1 = datetime(2024, 1, 10, 12, 0, 5, tzinfo=UTC) -def _workflow_run() -> WorkflowRun: - run = WorkflowRun(id="run-001", finished_at=_T1) - run.created_at = _T0 - return run +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] = { - "workflow_data": _workflow_run(), "node_executions": [], "conversation_id": None, "message_started_at": None, @@ -156,7 +165,7 @@ def make_draft_node_info(**overrides) -> DraftNodeExecutionTrace: def make_message_info(**overrides) -> MessageTraceInfo: defaults: dict[str, Any] = { "message_id": "msg-001", - "message_data": make_message(message_id="msg-001", app_id="app-001"), + "message_data": _trace_message(), "error": None, "file_list": [], "gen_ai_server_time_to_first_token": None, @@ -583,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..63ec791f20a1e7 --- /dev/null +++ b/api/tests/unit_tests/extensions/application_services/test_tracing.py @@ -0,0 +1,81 @@ +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 + + assert isinstance(agent_traces._traces, ApplicationTraceRepository) + 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 fdff1c556983b9..7d2124bc266ae5 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 @@ -154,10 +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.egress import SSRFEgress -from services.tracing.emitter import CeleryTraceEmitter -from services.tracing.provider_configs import ProviderConfigGateway +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, @@ -280,9 +278,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) @@ -290,28 +290,26 @@ 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( @@ -676,9 +674,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 - provider = services.app_tracing_configs._provider - assert isinstance(provider, ProviderConfigGateway) - assert isinstance(provider._egress, SSRFEgress) + 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_tracing.py b/api/tests/unit_tests/extensions/test_tracing.py index 78b68c29f881c8..64ee89e7f73309 100644 --- a/api/tests/unit_tests/extensions/test_tracing.py +++ b/api/tests/unit_tests/extensions/test_tracing.py @@ -1,19 +1,25 @@ -from collections.abc import Mapping +from typing import TYPE_CHECKING, Literal, override -import pytest +from dify_agent.layers.execution_context import ApplicationTracingContext from dify_trace_sdk.emitter import TraceScope -from dify_trace_sdk.events import TraceEvent +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, ) -from services.tracing.emitter import correlation_from_extras -from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks -from tests.helpers.unit.config_override import apply_config_overrides + +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: @@ -24,57 +30,131 @@ def emit(self, event: TraceEvent) -> None: self.events.append(event) -class _Scopes: +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 + ) - def __call__( + @override + def for_execution( self, *, tenant_id: str, app_id: str | None, app_tracing: str | None, user_id: str | None, - extras: Mapping[str, object], + correlation: TraceCorrelation, ) -> TraceScope | None: self.policies.append(app_tracing) - if self.fails: - raise RuntimeError("policy lookup failed") - return TraceScope( - emitter=_Recorder(), - tenant_id=tenant_id, - app_id=app_id, - user_id=user_id, - correlation=correlation_from_extras(extras), - ) + 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, extras: Mapping[str, object] + self, *, tenant_id: str, app_id: str | None, user_id: str | None, correlation: TraceCorrelation ) -> TraceScope | None: - return self(tenant_id=tenant_id, app_id=app_id, app_tracing=None, user_id=user_id, extras=extras) + 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 -def _app(*, fails: bool) -> Flask: + @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(fails=fails), hooks=RuntimeApplicationTracingHooks()) + bind_tracing(app, scopes=scopes, hooks=hooks) return app -def test_scopes_come_from_the_bound_factory() -> None: - with _app(fails=False).app_context(): - scope = create_trace_scope(tenant_id="tenant-1", app_id="app-1", app_tracing=None, user_id="u", extras={}) +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(fails=True).app_context(): + 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 ) @@ -83,14 +163,11 @@ def test_scopes_fail_open() -> None: assert create_trace_scope_for_app(tenant_id="tenant-1", app_id="app-1", user_id=None, extras={}) is None -def test_agent_hooks_follow_the_agent_tracing_switch(monkeypatch: pytest.MonkeyPatch) -> None: - app = _app(fails=False) +def test_agent_hooks_are_the_ones_the_composition_root_bound() -> None: + hooks = _Hooks() - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - with app.app_context(): + 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 - - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - with app.app_context(): - assert isinstance(application_tracing_hooks(), RuntimeApplicationTracingHooks) assert application_tracing_hooks() is None 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 06f3bcf80eece2..da8bd5da3655d8 100644 --- a/api/tests/unit_tests/repositories/test_application_trace_repository.py +++ b/api/tests/unit_tests/repositories/test_application_trace_repository.py @@ -1,14 +1,18 @@ import json from collections.abc import Callable, Iterator from dataclasses import dataclass -from datetime import timedelta +from datetime import UTC, datetime, timedelta +from decimal import Decimal from typing import Literal -from uuid import UUID +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 WorkflowRunCompleted +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace from dify_trace_sdk.provider import SendResult +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 @@ -18,27 +22,198 @@ 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, Message +from models.enums import ConversationFromSource +from models.model import App, AppMode, Message, TraceAppConfig from models.workflow import WorkflowNodeExecutionModel, WorkflowRun from repositories import application_trace_repository as repository_module -from repositories.application_trace_repository import TraceCollectionError -from services.tracing.records import ModerationTraceInfo -from services.tracing.recovery import _workflow_task +from repositories.application_trace_repository import ApplicationTraceRepository +from repositories.trace_record_repository import TraceRecordRepository +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import TraceInvocation, TraceRun, TraceRunState, TraceRunStatus +from services.tracing.errors import ( + TraceCollectionConflictError, + TraceCollectionError, + TraceCollectionNotFoundError, + TracingStoreUnavailableError, +) +from services.tracing.providers import ProviderConfigCodec +from services.tracing.records import ModerationTraceInfo, TraceMessage, WorkflowTraceInfo +from services.tracing.recovery import _workflow_event 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_no_egress as no_egress -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 ( - _application_trace_collection_tracing_service as tracing_service, +from tests.helpers.unit.model_factories import make_app, make_message, make_workflow +from tests.helpers.unit.services.tracing_ports import ( + AGENT_TRACE_SETTINGS, + InstalledTraceProviders, + WorkspaceBoundCipher, ) -from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection + +"""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](), +} + + +def identifier(value: str) -> str: + return str(uuid5(NAMESPACE_URL, value)) + + +@dataclass +class Collection: + repository: ApplicationTraceRepository + sessions: sessionmaker[Session] + run: TraceRun + + def invocation(self) -> TraceInvocation: + return self.repository.invocation( + run_id=self.run.id, agent_id=identifier("agent"), node_execution_id=None, attempt=0, 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( + {"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={ + "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 = repository.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 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) -> tuple[int, str]: + return collection.repository.ingest(batch, max_spans=100, max_bytes=1_000_000) + + +def terminal(collection: Collection) -> None: + 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,), + ) + collection.repository.terminal(run.id, trace.model_dump(mode="json"), root.span_id) + + +def claim(collection: Collection) -> tuple[TraceRun, str] | None: + return collection.repository.claim(collection.run.id, collection_timeout=60, lease_seconds=30) + + +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.""" @@ -49,22 +224,34 @@ class StoredSource: kind: Literal["workflow", "node"] row: dict[str, object] requests: list[str] + scheduled: list[str] - def start(self) -> ApplicationTraceRun: + def start(self) -> TraceRun: run = self.case.repository.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 tracing_service(self.case, no_egress()) + return ApplicationTracingService( + traces=self.case.repository, + records=TraceRecordRepository(self.case.sessions, files_url="https://files.test", enterprise_details=False), + builder=CanonicalTraceBuilder(), + # No provider is installed, so a delivery stops before it reaches any destination. + providers=InstalledTraceProviders(), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + egress=RecordingEgress([], grpc_target=None), + schedule_export=self.scheduled.append, + settings=AGENT_TRACE_SETTINGS, + ) - 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, @@ -182,16 +369,14 @@ def execute_sql( event.listen(sqlite_engine, "commit", ended) event.listen(sqlite_engine, "rollback", ended) try: - yield StoredSource(case, kind, row, requests) + yield StoredSource(case, kind, row, requests, scheduled=[]) finally: event.remove(sqlite_engine, "begin", begun) event.remove(sqlite_engine, "commit", ended) 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 @@ -199,23 +384,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.scheduled == [run.id] assert len(stored.requests) == 1 @staticmethod @@ -223,11 +406,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 @@ -238,8 +424,8 @@ 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 @@ -251,7 +437,7 @@ def test_initialization_failure_uses_configured_source_reader(stored: StoredSour 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 @@ -265,6 +451,7 @@ def test_existing_snapshot_reentry_does_not_consult_logstore(stored: StoredSourc 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 @@ -287,20 +474,22 @@ def test_workflow_node_deep_link_is_scoped_to_durable_invocation(stored: StoredS case.repository.complete( run.id, lease[1], - state="accepted", + state=TraceRunState.ACCEPTED, + error_code=None, result=SendResult( trace_id="provider-id", trace_url=None, span_urls={invocation.span_id: "https://trace.test/span"} ), + retry_delay=None, ) - 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")) + 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": - event = _workflow_task(stored.case.repository, run, naive_utc_now()) + event = _workflow_event(stored.case.repository, run, naive_utc_now()) assert event is not None assert event.payload == WorkflowRunCompleted( workflow_run_id=run.source_id, conversation_id=None, total_tokens=None @@ -309,9 +498,10 @@ def test_preview_settlement_and_workflow_recovery_use_configured_source(stored: assert event.user_id == stored.row["created_by"] assert event.correlation.session_id == run.session_id stored.row["status"] = "running" - assert _workflow_task(stored.case.repository, run, naive_utc_now()) is None + assert _workflow_event(stored.case.repository, run, naive_utc_now()) is None 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 @@ -325,10 +515,10 @@ def test_preview_settlement_and_workflow_recovery_use_configured_source(stored: @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), None) - assert stored.service().capture_terminal(info) is True - result = stored.status() - assert result["status"] == "failed" - assert result["error_code"] == "tracing_initialization_failed" + 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) @@ -349,7 +539,7 @@ 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) @@ -367,7 +557,20 @@ def test_auxiliary_content_is_suppressed_for_logstore_chatflow(stored: StoredSou stored.start() info = ModerationTraceInfo( message_id=identifier("workflow-log"), - message_data=message, + 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", @@ -376,15 +579,16 @@ def test_auxiliary_content_is_suppressed_for_logstore_chatflow(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) @@ -396,7 +600,7 @@ def test_linked_parent_requires_logstore_run_and_durable_owner_to_agree( info = workflow_info(_dict_to_workflow_run(stored.row), _dict_to_workflow_run(stored.row)) if mismatch: stored.row[mismatch] = identifier("foreign") - linked = stored.service().link_managed_parent(info) + linked = stored.service().link_managed_parent(tenant_id=stored.case.run.tenant_id, info=info) if mismatch: assert linked is info else: @@ -420,14 +624,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 ) @@ -463,13 +667,14 @@ def run_count(case: Collection) -> int | None: ) -def start_message(case: Collection, session_id: str) -> ApplicationTraceRun | None: +def start_message(case: Collection, session_id: str) -> TraceRun | None: return case.repository.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, ) @@ -480,14 +685,14 @@ def test_concurrent_start_returns_the_first_frozen_snapshot( ) -> None: case = seed(sqlite_session_factory, "concurrent-start") drop_run(case) - winners: list[ApplicationTraceRun | None] = [] + 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 == "collecting" + assert current.state is TraceRunState.COLLECTING assert run_count(case) == 1 @staticmethod @@ -496,7 +701,7 @@ def test_initialization_failure_yields_to_a_concurrently_frozen_start( ) -> None: case = seed(sqlite_session_factory, "concurrent-failure") drop_run(case) - winners: list[ApplicationTraceRun | None] = [] + winners: list[TraceRun | None] = [] race_before_insert(monkeypatch, lambda: winners.append(start_message(case, "started-session"))) current = 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 @@ -504,7 +709,7 @@ def test_initialization_failure_yields_to_a_concurrently_frozen_start( assert current is not None assert winners[0] is not None assert current.id == winners[0].id - assert current.state == "collecting" + 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 @@ -517,8 +722,9 @@ def test_identifier_collision_is_not_mistaken_for_a_concurrent_start( other = seed(sqlite_session_factory, "colliding-owner") drop_run(case) monkeypatch.setattr(repository_module, "uuidv7", lambda: UUID(other.run.id)) - with pytest.raises(IntegrityError): + 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 @@ -539,9 +745,10 @@ def test_malformed_policy_freezes_disabled_decision( source_kind="message", source_id=case.run.source_id, session_id="", + workflow_id=None, ) 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 @@ -575,7 +782,7 @@ def test_initialization_failure_cannot_replace_snapshot_or_create_foreign_row( ) 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 assert ( case.repository.record_initialization_failure( @@ -595,22 +802,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 @@ -639,7 +847,7 @@ 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" + assert failed.state is TraceRunState.FAILED case.repository.terminal(case.run.id, {"replacement": True}, "f" * 16) frozen = case.repository.get(case.run.id) assert frozen is not None @@ -647,15 +855,18 @@ def test_first_terminal_is_immutable_and_revoked_owner_is_never_exported( @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) @@ -663,20 +874,20 @@ 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): + with pytest.raises(error, match=code): 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 + 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"): + with pytest.raises(TraceCollectionNotFoundError, 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"): + with pytest.raises(TraceCollectionConflictError, match="tracing_run_not_collecting"): case.repository.invocation( run_id=identifier("missing"), agent_id=identifier("agent"), @@ -694,7 +905,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( @@ -708,10 +919,10 @@ 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 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 expired.repository.purge(retention_days=7, limit=100) == 0 diff --git a/api/tests/unit_tests/services/tracing/test_assembly.py b/api/tests/unit_tests/repositories/test_trace_record_repository.py similarity index 81% rename from api/tests/unit_tests/services/tracing/test_assembly.py rename to api/tests/unit_tests/repositories/test_trace_record_repository.py index 4c290f582d7b33..51f882d3fdaa62 100644 --- a/api/tests/unit_tests/services/tracing/test_assembly.py +++ b/api/tests/unit_tests/repositories/test_trace_record_repository.py @@ -10,12 +10,14 @@ TraceEvent, TracePayload, ) +from sqlalchemy import create_engine from sqlalchemy.orm import Session, sessionmaker from models.enums import ConversationFromSource, ConversationStatus, EndUserType, InvokeFrom from models.model import App, AppMode, Conversation, EndUser, Message -from services.tracing.assembly import SourceNotFoundError, SourceNotReadyError, TraceAssembler -from services.tracing.records import MessageTraceInfo, ModerationTraceInfo, ToolTraceInfo +from repositories.trace_record_repository import TraceRecordRepository +from services.tracing.errors import SourceNotFoundError, SourceNotReadyError, TracingStoreUnavailableError +from services.tracing.records import MessageTraceInfo, ModerationTraceInfo, ToolTraceInfo, TraceMessage TENANT_ID = "tenant-1" APP_ID = "app-1" @@ -89,8 +91,8 @@ def records(sqlite_session: Session) -> None: sqlite_session.commit() -def _assembler(sessions: sessionmaker[Session]) -> TraceAssembler: - return TraceAssembler(sessions, files_url="https://files.example.com", enterprise_details=False) +def _assembler(sessions: sessionmaker[Session]) -> TraceRecordRepository: + return TraceRecordRepository(sessions, files_url="https://files.example.com", enterprise_details=False) def _event(payload: TracePayload, *, tenant_id: str = TENANT_ID) -> TraceEvent: @@ -203,3 +205,33 @@ def test_events_without_a_message_id_are_not_found(sqlite_session_factory: sessi 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"))) 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 2f7078091617a0..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 @@ -13,7 +13,7 @@ from services.agent_tracing_config_service import AgentTracingConfigService from services.app.agent_app_contracts import AgentAppNotFoundError from services.app_tracing_config_service import AppTraceSettings, AppTracingConfigService -from services.tracing.provider_configs import ProviderConfigGateway +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, ...]] @@ -57,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=ProviderConfigGateway(RecordingEgress([], grpc_target=None)), + provider=deployment_provider_configs(RecordingEgress([], grpc_target=None)), ) calls: list[RecordedCall] = [] 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 0df12fc5f4596d..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 @@ -23,6 +23,7 @@ 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.""" @@ -63,7 +64,7 @@ def case(sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.Monk ), ) ) - gateway = ProviderConfigGateway(RecordingEgress([], grpc_target=None)) + gateway = deployment_provider_configs(RecordingEgress([], grpc_target=None)) def prepare(*, tracing_config: dict[str, object], **_: object) -> dict[str, object]: return dict(tracing_config) 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 43413d13bd4f66..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()), @@ -194,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()), @@ -228,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 a3f3ede270cde4..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 @@ -25,10 +25,10 @@ from core.plugin.impl.base import _get_plugin_daemon_request_timeout from extensions.ext_database import db from extensions.ext_redis import RedisClientWrapper -from extensions.tracing import bind_tracing 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, @@ -36,10 +36,10 @@ _default_suggested_questions_model_parameters, ) from services.message_suggested_questions_service import SuggestedQuestionsContext -from services.tracing.emitter import TraceScopes -from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks +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 @@ -85,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() @@ -166,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, @@ -248,10 +250,12 @@ def record_session(session: Session, _transaction: SessionTransaction, _connecti factory = db.session.session_factory recorder = _TraceRecorder() - bind_tracing( - current_app, - scopes=TraceScopes(emitter=recorder, session_factory=sqlite_session_factory, enterprise_enabled=lambda: False), - hooks=RuntimeApplicationTracingHooks(), + 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) @@ -259,7 +263,7 @@ def record_session(session: Session, _transaction: SessionTransaction, _connecti # Empty detached history needs no token-counting call. The real failed # model lookup returns [], then the trace scope reads the stored policy. with current_app.app_context(): - questions = SuggestedQuestionsGenerator().generate( + questions = generator.generate( PreparedHistory(prompts=()), history_model=make_model_instance(provider="invalid/provider", model="unused"), context=context, @@ -317,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_persistence.py b/api/tests/unit_tests/services/test_workflow_persistence.py index fa6aa73d73d59a..ea38df25573781 100644 --- a/api/tests/unit_tests/services/test_workflow_persistence.py +++ b/api/tests/unit_tests/services/test_workflow_persistence.py @@ -3,7 +3,7 @@ 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 @@ -73,10 +73,8 @@ DifyAPISQLAlchemyWorkflowNodeExecutionRepository, ) from services import workflow_persistence as persistence_mod -from services.tracing.emitter import TraceScopes -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. @@ -400,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}) @@ -901,44 +894,43 @@ 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)) app = Flask(__name__) bind_tracing( app, - scopes=TraceScopes( - emitter=_TraceRecorder(workflow_repo), - session_factory=sessionmaker(), - enterprise_enabled=lambda: False, + scopes=kernel_trace_scopes( + _TraceRecorder(workflow_repo), policies=NoTracingPolicies(), enterprise_enabled=False ), - hooks=RuntimeApplicationTracingHooks(), + hooks=_FreezeAfterSave(None), ) with app.app_context(): 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..5afba1f01b2a70 --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_adapters.py @@ -0,0 +1,205 @@ +from datetime import UTC, datetime +from typing import override + +import pytest +from dify_trace_sdk.events import ModerationChecked, TraceCorrelation, TraceEvent, WorkflowRunCompleted +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace + +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, + RETRY_TASK, + TRACE_QUEUE, + 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 _Broker: + def __init__(self) -> None: + self.sent: list[tuple[str, tuple[object, ...], str]] = [] + self.countdowns: list[int] = [] + + def send_task(self, name: str, *, args: tuple[object, ...], queue: str, countdown: int = 0) -> None: + self.sent.append((name, args, queue)) + self.countdowns.append(countdown) + + +class _Storage: + def __init__(self) -> None: + self.files: dict[str, bytes] = {} + + def save(self, key: str, data: bytes) -> None: + self.files[key] = data + + +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(monkeypatch: pytest.MonkeyPatch) -> _Broker: + broker = _Broker() + monkeypatch.setattr(adapters, "celery_app", broker) + return broker + + +@pytest.fixture +def storage(monkeypatch: pytest.MonkeyPatch) -> _Storage: + storage = _Storage() + monkeypatch.setattr(adapters, "storage", storage) + return storage + + +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: _Broker, storage: _Storage) -> None: + event = _event("hello") + + CeleryTraceEmitter().emit(event) + + ((name, (envelope,), queue),) = broker.sent + assert (name, queue) == (DELIVER_TASK, TRACE_QUEUE) + assert isinstance(envelope, dict) + assert TraceEvent.model_validate_json(envelope["event"]) == event + assert storage.files == {} + + +def test_large_events_travel_through_storage(broker: _Broker, storage: _Storage) -> None: + event = _event("x" * INLINE_EVENT_BYTES) + + CeleryTraceEmitter().emit(event) + + ((_, (envelope,), _),) = broker.sent + 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.files["trace_payloads/2026-03-04/events/event-1.json"]) == event + + +@pytest.mark.usefixtures("storage") +def test_inline_size_is_measured_in_bytes(broker: _Broker) -> None: + CeleryTraceEmitter().emit(_event("中" * (INLINE_EVENT_BYTES // 2))) + + ((_, (envelope,), _),) = broker.sent + assert isinstance(envelope, dict) + assert "payload_key" in envelope + + +@pytest.mark.usefixtures("storage") +def test_workflow_runs_wait_for_node_executions_to_land(broker: _Broker) -> None: + workflow = _event("hello").model_copy( + update={"payload": WorkflowRunCompleted(workflow_run_id="run-1", conversation_id=None, total_tokens=0)} + ) + + CeleryTraceEmitter().emit(_event("hello")) + CeleryTraceEmitter().emit(workflow) + + assert broker.countdowns == [0, WORKFLOW_SETTLE_SECONDS] + + +def test_export_retries_persist_the_trace_and_run_after_their_delay(broker: _Broker, storage: _Storage) -> 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) + + ((name, (key,), queue),) = broker.sent + assert (name, queue, broker.countdowns) == (RETRY_TASK, TRACE_QUEUE, [20]) + assert isinstance(key, str) + assert key.startswith("trace_payloads/") + assert "/exports/" in key + assert ExportJob.model_validate_json(storage.files[key]) == job + + +def test_agent_trace_exports_run_on_the_trace_queue(broker: _Broker) -> None: + schedule_agent_trace_export("run-1") + + assert broker.sent == [(AGENT_TRACE_EXPORT_TASK, ("run-1",), TRACE_QUEUE)] + + +@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 index 6b872cce42b5bc..1c22b6f305b150 100644 --- a/api/tests/unit_tests/services/tracing/test_canonical.py +++ b/api/tests/unit_tests/services/tracing/test_canonical.py @@ -6,13 +6,12 @@ from graphon.entities import WorkflowNodeExecution from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus -from models.model import Message -from models.workflow import WorkflowRun from services.tracing.canonical import CanonicalTraceBuilder, message_trace_id from services.tracing.records import ( MessageTraceInfo, PromptGenerationTraceInfo, ToolTraceInfo, + TraceMessage, WorkflowTraceInfo, ) @@ -22,21 +21,21 @@ WORKFLOW_RUN_ID = "4b4b8c2e-1a0f-4bd0-93c1-0f2f0e9d6a33" -def _message() -> Message: - message = Message( +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", - error=None, ) - message.created_at = STARTED_AT - return message def _message_record(**metadata: object) -> MessageTraceInfo: @@ -80,12 +79,6 @@ def _node(execution_id: str, node_type: str, *, index: int, offset: int, **field ) -def _workflow_run() -> WorkflowRun: - run = WorkflowRun(id=WORKFLOW_RUN_ID, finished_at=STARTED_AT + timedelta(seconds=3)) - run.created_at = STARTED_AT - return run - - def _workflow_record(*, message_id: str | None, metadata: dict[str, object]) -> WorkflowTraceInfo: llm = _node( "exec-llm", @@ -104,7 +97,6 @@ def _workflow_record(*, message_id: str | None, metadata: dict[str, object]) -> start = _node("exec-start", BuiltinNodeTypes.START, index=1, offset=0, inputs={"query": "hi"}) return WorkflowTraceInfo( message_id=message_id, - workflow_data=_workflow_run(), 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, diff --git a/api/tests/unit_tests/services/tracing/test_delivery.py b/api/tests/unit_tests/services/tracing/test_delivery.py index caf05c5b0ab3d2..3b80818652068d 100644 --- a/api/tests/unit_tests/services/tracing/test_delivery.py +++ b/api/tests/unit_tests/services/tracing/test_delivery.py @@ -1,31 +1,37 @@ -from collections.abc import Callable +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 +from dify_trace_sdk.testing import RecordingEgress, response from pydantic import ValidationError -from sqlalchemy.orm import sessionmaker -from models.model import Message -from services.tracing.assembly import SourceNotReadyError from services.tracing.canonical import CanonicalTraceBuilder from services.tracing.delivery import ( - Destination, ExportJob, - TraceCollectorUnavailableError, + ExportRetryQueue, + StoredTraceDestination, + TraceCollector, TraceDeliveryService, + TraceDestinationQuery, + TraceRecordSource, retry_delay_seconds, ) -from services.tracing.records import BaseTraceInfo, MessageTraceInfo, WorkflowTraceInfo +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): @@ -33,15 +39,12 @@ class _Config(ProviderConfig): class _Provider(TraceProvider[_Config]): - id = "scripted" + """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) - def __init__(self, outcomes: list[Exception | None]) -> None: - super().__init__(RecordingEgress([], grpc_target=None)) - self.outcomes = outcomes - self.sent: list[Trace] = [] - @override def verify(self, config: _Config) -> _Config: return config @@ -52,18 +55,23 @@ def project_url(self, config: _Config) -> str | None: @override def send(self, config: _Config, trace: Trace) -> SendResult: - self.sent.append(trace) - outcome = self.outcomes.pop(0) - if outcome is not None: - raise outcome + 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 _Assembler: +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): @@ -71,59 +79,77 @@ def assemble(self, event: TraceEvent) -> BaseTraceInfo: return self.result -class _Collector: +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[BaseTraceInfo] = [] + self.seen: list[tuple[str, BaseTraceInfo]] = [] - def capture_terminal(self, info: BaseTraceInfo) -> bool: + @override + def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: if self.fails: raise RuntimeError("ledger unavailable") - self.seen.append(info) + self.seen.append((tenant_id, info)) return self.captures - def link_managed_parent(self, info: WorkflowTraceInfo) -> WorkflowTraceInfo: + @override + def link_managed_parent(self, *, tenant_id: str, info: WorkflowTraceInfo) -> WorkflowTraceInfo: return info -class _Service(TraceDeliveryService): - """Delivery with an in-memory destination and retry schedule.""" - - def __init__( - self, - *, - assembler: _Assembler, - destination: Destination | None, - enterprise_sink: Callable[[BaseTraceInfo], None] | None, - collector: _Collector | None, - ) -> None: - super().__init__( - session_factory=sessionmaker(), - assembler=assembler, +@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=RecordingEgress([], grpc_target=None), - enterprise_sink=enterprise_sink, + egress=self.egress, + retries=self.retries, + enterprise_sink=self.enterprise.append if enterprise else None, collector=collector, max_export_attempts=3, ) - self._destination = destination - self.scheduled: list[ExportJob] = [] - - @override - def destination(self, tenant_id: str, app_id: str) -> Destination | None: - return self._destination - @override - def _schedule_retry(self, job: ExportJob) -> None: - self.scheduled.append(job) + 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-1", - app_id="app-1", + tenant_id=TENANT_ID, + app_id=APP_ID, user_id="account-1", occurred_at=datetime.now(UTC), correlation=TraceCorrelation( @@ -133,15 +159,21 @@ def _event() -> TraceEvent: ) -def _record(*, workflow_run_id: str | None = None) -> MessageTraceInfo: - message = Message( +def _record(*, workflow_run_id: str | None) -> MessageTraceInfo: + message = TraceMessage( id=MESSAGE_ID, - app_id="app-1", + app_id=APP_ID, conversation_id="conversation-1", - answer="hello", 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", ) - message.created_at = STARTED_AT return MessageTraceInfo( message_id=MESSAGE_ID, message_data=message, @@ -163,139 +195,135 @@ def _record(*, workflow_run_id: str | None = None) -> MessageTraceInfo: ) -def _destination(outcomes: list[Exception | None]) -> Destination: - return Destination(provider=_Provider(outcomes), config=_Config(endpoint="https://collector.example")) +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: - destination = _destination([None]) - enterprise: list[BaseTraceInfo] = [] + delivery = _delivery(_record(workflow_run_id=None), [response(200, {})]) collector = _Collector(captures=False, fails=False) - service = _Service( - assembler=_Assembler(_record()), destination=destination, enterprise_sink=enterprise.append, collector=collector - ) - service.deliver(_event(), enterprise_delivered=False) + delivery.service(enterprise=True, collector=collector).deliver(_event(), enterprise_delivered=False) - assert len(enterprise) == 1 - assert len(collector.seen) == 1 - assert isinstance(destination.provider, _Provider) - (trace,) = destination.provider.sent + 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 service.scheduled == [] + assert delivery.destinations.lookups == [(TENANT_ID, APP_ID)] + assert delivery.retries.scheduled == [] def test_enterprise_sink_is_skipped_on_redelivery() -> None: - enterprise: list[BaseTraceInfo] = [] - service = _Service( - assembler=_Assembler(_record()), - destination=_destination([None]), - enterprise_sink=enterprise.append, - collector=None, - ) + delivery = _delivery(_record(workflow_run_id=None), [response(200, {})]) - service.deliver(_event(), enterprise_delivered=True) + delivery.service(enterprise=True, collector=None).deliver(_event(), enterprise_delivered=True) - assert enterprise == [] + assert delivery.enterprise == [] def test_records_the_collector_owns_never_reach_the_provider() -> None: - destination = _destination([]) - service = _Service( - assembler=_Assembler(_record()), - destination=destination, - enterprise_sink=None, - collector=_Collector(captures=True, fails=False), - ) + delivery = _delivery(_record(workflow_run_id=None), []) - service.deliver(_event(), enterprise_delivered=False) + delivery.service(enterprise=False, collector=_Collector(captures=True, fails=False)).deliver( + _event(), enterprise_delivered=False + ) - assert isinstance(destination.provider, _Provider) - assert destination.provider.sent == [] + assert delivery.sent() == [] + assert delivery.destinations.lookups == [] def test_an_unavailable_collector_asks_for_redelivery() -> None: - service = _Service( - assembler=_Assembler(_record()), - destination=_destination([]), - enterprise_sink=None, - collector=_Collector(captures=False, fails=True), - ) + delivery = _delivery(_record(workflow_run_id=None), []) with pytest.raises(TraceCollectorUnavailableError): - service.deliver(_event(), enterprise_delivered=False) + 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: - destination = _destination([]) - service = _Service( - assembler=_Assembler(_record(workflow_run_id="run-1")), - destination=destination, - enterprise_sink=None, - collector=None, - ) + delivery = _delivery(_record(workflow_run_id="run-1"), []) - service.deliver(_event(), enterprise_delivered=False) + delivery.service(enterprise=False, collector=None).deliver(_event(), enterprise_delivered=False) - assert isinstance(destination.provider, _Provider) - assert destination.provider.sent == [] + assert delivery.sent() == [] def test_source_not_ready_propagates_for_redelivery() -> None: - service = _Service( - assembler=_Assembler(SourceNotReadyError()), destination=_destination([]), enterprise_sink=None, collector=None - ) + delivery = _delivery(SourceNotReadyError(), []) with pytest.raises(SourceNotReadyError): - service.deliver(_event(), enterprise_delivered=False) + 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: - service = _Service( - assembler=_Assembler(_record()), - destination=_destination([ExportRetryableError("busy"), ExportRetryableError("busy")]), - enterprise_sink=None, - collector=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,) = service.scheduled - assert (job.attempt, job.provider_id, job.app_id) == (2, "scripted", "app-1") - exhausted = job.model_copy(update={"attempt": 3}) - destination = service.destination("tenant-1", "app-1") + ((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(exhausted, destination) - assert len(service.scheduled) == 1 + 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: - service = _Service( - assembler=_Assembler(_record()), destination=_destination([failure]), enterprise_sink=None, collector=None - ) + delivery = _delivery(_record(workflow_run_id=None), [failure]) - service.deliver(_event(), enterprise_delivered=False) + delivery.service(enterprise=False, collector=None).deliver(_event(), enterprise_delivered=False) - assert service.scheduled == [] + assert delivery.retries.scheduled == [] def test_retry_is_dropped_when_the_app_now_traces_elsewhere() -> None: - trace = CanonicalTraceBuilder().build(_record()) + 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 - destination = _destination([]) - service = _Service(assembler=_Assembler(_record()), destination=destination, enterprise_sink=None, collector=None) + delivery = _delivery(_record(workflow_run_id=None), [response(200, {})]) - service.retry(ExportJob(tenant_id="tenant-1", app_id="app-1", provider_id="langfuse", attempt=2, trace=trace)) + delivery.service(enterprise=False, collector=None).retry( + ExportJob(tenant_id=TENANT_ID, app_id=APP_ID, provider_id="langfuse", attempt=2, trace=trace) + ) - assert isinstance(destination.provider, _Provider) - assert destination.provider.sent == [] + 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()) + trace = CanonicalTraceBuilder().build(_record(workflow_run_id=None)) assert trace is not None - job = ExportJob(tenant_id="tenant-1", app_id="app-1", provider_id="scripted", attempt=2, trace=trace) + 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 @@ -307,25 +335,33 @@ def test_retry_delay_grows_and_is_capped() -> None: def test_records_that_fail_validation_are_dropped_instead_of_failing_the_task() -> None: with pytest.raises(ValidationError) as invalid: MessageTraceInfo.model_validate({}) - destination = _destination([]) - service = _Service( - assembler=_Assembler(invalid.value), destination=destination, enterprise_sink=None, collector=None - ) + delivery = _delivery(invalid.value, []) - service.deliver(_event(), enterprise_delivered=False) + delivery.service(enterprise=False, collector=None).deliver(_event(), enterprise_delivered=False) - assert isinstance(destination.provider, _Provider) - assert destination.provider.sent == [] + assert delivery.sent() == [] def test_unclassified_provider_failures_are_contained() -> None: - service = _Service( - assembler=_Assembler(_record()), - destination=_destination([RuntimeError("provider bug with secret-detail")]), - enterprise_sink=None, - collector=None, - ) + delivery = _delivery(_record(workflow_run_id=None), [RuntimeError("provider bug with secret-detail")]) - service.deliver(_event(), enterprise_delivered=False) + 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 service.scheduled == [] + assert delivery.sent() == [] diff --git a/api/tests/unit_tests/services/tracing/test_emitter.py b/api/tests/unit_tests/services/tracing/test_emitter.py deleted file mode 100644 index cd9c464b8cb18d..00000000000000 --- a/api/tests/unit_tests/services/tracing/test_emitter.py +++ /dev/null @@ -1,202 +0,0 @@ -from datetime import UTC, datetime - -import pytest -from dify_trace_sdk.events import ModerationChecked, TraceCorrelation, TraceEvent, WorkflowRunCompleted -from sqlalchemy.orm import sessionmaker - -from core.helper.trace_id_helper import ParentTraceContext -from services.tracing import emitter -from services.tracing.emitter import ( - DELIVER_TASK, - INLINE_EVENT_BYTES, - TRACE_QUEUE, - WORKFLOW_SETTLE_SECONDS, - CeleryTraceEmitter, - TraceScopes, - correlation_from_extras, - payload_key, - tracing_policy_enabled, -) - - -class _Broker: - def __init__(self) -> None: - self.sent: list[tuple[str, tuple[dict[str, str]], str]] = [] - self.countdowns: list[int] = [] - - def send_task(self, name: str, *, args: tuple[dict[str, str]], queue: str, countdown: int) -> None: - self.sent.append((name, args, queue)) - self.countdowns.append(countdown) - - -class _Storage: - def __init__(self) -> None: - self.files: dict[str, bytes] = {} - - def save(self, key: str, data: bytes) -> None: - self.files[key] = data - - -class _Recorder: - def __init__(self) -> None: - self.events: list[TraceEvent] = [] - - def emit(self, event: TraceEvent) -> None: - self.events.append(event) - - -def _event(query: str) -> TraceEvent: - started_at = datetime(2026, 3, 4, 5, 6, tzinfo=UTC) - return TraceEvent( - event_id="event-1", - tenant_id="tenant-1", - app_id="app-1", - user_id=None, - occurred_at=started_at, - correlation=TraceCorrelation( - external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None - ), - payload=ModerationChecked( - message_id="message-1", - inputs={}, - query=query, - flagged=False, - action="direct_output", - preset_response="", - started_at=started_at, - ended_at=started_at, - ), - ) - - -@pytest.mark.parametrize( - ("policy", "enabled"), - [ - ('{"enabled": true, "tracing_provider": "langfuse"}', True), - ('{"enabled": false, "tracing_provider": "langfuse"}', False), - ('{"enabled": true, "tracing_provider": null}', False), - ("not json", False), - (None, False), - ], -) -def test_tracing_policy(policy: str | None, enabled: bool) -> None: - assert tracing_policy_enabled(policy) is enabled - - -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_small_events_travel_inline(monkeypatch: pytest.MonkeyPatch) -> None: - broker, storage = _Broker(), _Storage() - monkeypatch.setattr(emitter, "celery_app", broker) - monkeypatch.setattr(emitter, "storage", storage) - event = _event("hello") - - CeleryTraceEmitter().emit(event) - - ((name, (envelope,), queue),) = broker.sent - assert (name, queue) == (DELIVER_TASK, TRACE_QUEUE) - assert TraceEvent.model_validate_json(envelope["event"]) == event - assert storage.files == {} - - -def test_large_events_travel_through_storage(monkeypatch: pytest.MonkeyPatch) -> None: - broker, storage = _Broker(), _Storage() - monkeypatch.setattr(emitter, "celery_app", broker) - monkeypatch.setattr(emitter, "storage", storage) - event = _event("x" * INLINE_EVENT_BYTES) - - CeleryTraceEmitter().emit(event) - - ((_, (envelope,), _),) = broker.sent - assert envelope == {"payload_key": payload_key(event.occurred_at, "events", "event-1")} - assert envelope["payload_key"] == "trace_payloads/2026-03-04/events/event-1.json" - assert TraceEvent.model_validate_json(storage.files[envelope["payload_key"]]) == event - - -def _scopes(*, enterprise: bool) -> tuple[TraceScopes, _Recorder]: - recorder = _Recorder() - return TraceScopes( - emitter=recorder, session_factory=sessionmaker(), enterprise_enabled=lambda: enterprise - ), recorder - - -def test_no_scope_exists_when_nothing_consumes_the_events() -> None: - scopes, _ = _scopes(enterprise=False) - - assert scopes(tenant_id="tenant-1", app_id="app-1", app_tracing=None, user_id=None, extras={}) is None - - -@pytest.mark.parametrize( - ("policy", "enterprise"), - [('{"enabled": true, "tracing_provider": "langfuse"}', False), (None, True)], -) -def test_scopes_stamp_owner_and_correlation(policy: str | None, enterprise: bool) -> None: - scopes, recorder = _scopes(enterprise=enterprise) - - scope = scopes( - tenant_id="tenant-1", - app_id="app-1", - app_tracing=policy, - user_id="end-user-session", - extras={"external_trace_id": "caller-trace"}, - ) - - 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.external_trace_id == "caller-trace" - assert scope.emitter is recorder - - -def test_app_scopes_without_an_app_depend_on_enterprise_telemetry() -> None: - scopes, _ = _scopes(enterprise=True) - - scope = scopes.for_app(tenant_id="tenant-1", app_id=None, user_id=None, extras={}) - - assert scope is not None - assert scope.app_id is None - - -def test_workflow_runs_wait_for_node_executions_to_land(monkeypatch: pytest.MonkeyPatch) -> None: - broker = _Broker() - monkeypatch.setattr(emitter, "celery_app", broker) - workflow = _event("hello").model_copy( - update={"payload": WorkflowRunCompleted(workflow_run_id="run-1", conversation_id=None, total_tokens=0)} - ) - - CeleryTraceEmitter().emit(_event("hello")) - CeleryTraceEmitter().emit(workflow) - - assert broker.countdowns == [0, WORKFLOW_SETTLE_SECONDS] - - -def test_inline_size_is_measured_in_bytes(monkeypatch: pytest.MonkeyPatch) -> None: - broker, storage = _Broker(), _Storage() - monkeypatch.setattr(emitter, "celery_app", broker) - monkeypatch.setattr(emitter, "storage", storage) - - CeleryTraceEmitter().emit(_event("\u4e2d" * (INLINE_EVENT_BYTES // 2))) - - ((_, (envelope,), _),) = broker.sent - assert "payload_key" in envelope diff --git a/api/tests/unit_tests/services/tracing/test_provider_configs.py b/api/tests/unit_tests/services/tracing/test_provider_configs.py index 9a5ea8ed28df60..6ba64cc23ef4cf 100644 --- a/api/tests/unit_tests/services/tracing/test_provider_configs.py +++ b/api/tests/unit_tests/services/tracing/test_provider_configs.py @@ -1,43 +1,35 @@ 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 dify_trace_sdk.egress import EgressResponse +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 import provider_configs, registry from services.tracing.provider_configs import ( ProviderConfigGateway, - ProviderConfigUnreadableError, - capabilities_payload, - configured_secrets, - decode_config, - encode_config, + TraceProviderCapabilitiesQuery, + TraceProviderSupport, ) +from services.tracing.providers import ProviderConfigCodec, 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 f"enc:{tenant_id}:{token}" + return CIPHER.encrypt(tenant_id, token) -def _batch_decrypt(tenant_id: str, tokens: list[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") - return [token.removeprefix(prefix) for token in tokens] - - -@pytest.fixture(autouse=True) -def reversible_encryption(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(provider_configs, "encrypt_token", _encrypt) - monkeypatch.setattr(provider_configs, "batch_decrypt_token", _batch_decrypt) +def _gateway(egress: Egress) -> ProviderConfigGateway: + return ProviderConfigGateway( + providers=InstalledTraceProviders(LangfuseProvider, PhoenixProvider), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + egress=egress, + ) def _projects(project_id: str) -> EgressResponse: @@ -54,23 +46,8 @@ def _stored(**overrides: object) -> dict[str, object]: } -def test_stored_configs_encrypt_only_secrets_and_round_trip() -> None: - config = LangfuseConfig(public_key=PUBLIC_KEY, secret_key=SECRET_KEY, host="https://langfuse.example.com") - - stored = encode_config(WORKSPACE, config) - - assert stored["public_key"] == _encrypt(WORKSPACE, PUBLIC_KEY) - assert stored["host"] == "https://langfuse.example.com" - assert decode_config(WORKSPACE, LangfuseProvider, stored) == config - - -def test_a_config_encrypted_for_another_workspace_is_unreadable() -> None: - with pytest.raises(ProviderConfigUnreadableError): - decode_config("tenant-2", LangfuseProvider, _stored()) - - def test_presentation_never_returns_secrets_or_derived_fields() -> None: - gateway = ProviderConfigGateway(RecordingEgress([], grpc_target=None)) + gateway = _gateway(RecordingEgress([], grpc_target=None)) presented = gateway.present_config(workspace_id=WORKSPACE, tracing_provider="langfuse", tracing_config=_stored()) @@ -84,7 +61,7 @@ def test_presentation_never_returns_secrets_or_derived_fields() -> None: def test_new_configs_are_verified_and_store_derived_fields() -> None: egress = RecordingEgress([_projects("project-9")], grpc_target=None) - gateway = ProviderConfigGateway(egress) + gateway = _gateway(egress) stored = gateway.prepare_new_config( workspace_id=WORKSPACE, @@ -109,7 +86,7 @@ def test_new_configs_are_verified_and_store_derived_fields() -> None: ids=["omitted", "blank", "masked"], ) def test_updates_keep_a_secret_that_was_not_re_entered(submitted_secret: str | None) -> None: - gateway = ProviderConfigGateway(RecordingEgress([_projects("project-1")], grpc_target=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 @@ -125,7 +102,7 @@ def test_updates_keep_a_secret_that_was_not_re_entered(submitted_secret: str | N def test_updates_replace_a_re_entered_secret() -> None: - gateway = ProviderConfigGateway(RecordingEgress([_projects("project-1")], grpc_target=None)) + gateway = _gateway(RecordingEgress([_projects("project-1")], grpc_target=None)) stored = gateway.prepare_updated_config( workspace_id=WORKSPACE, @@ -142,7 +119,7 @@ def test_updates_replace_a_re_entered_secret() -> None: def test_failed_verification_is_reported_without_details() -> None: - gateway = ProviderConfigGateway(RecordingEgress([response(401, {"message": "bad key"})], grpc_target=None)) + gateway = _gateway(RecordingEgress([response(401, {"message": "bad key"})], grpc_target=None)) with pytest.raises(config_service.AppTracingConfigVerificationFailedError): gateway.prepare_new_config( @@ -156,7 +133,7 @@ def test_invalid_configs_are_rejected_before_any_request() -> None: egress = RecordingEgress([], grpc_target=None) with pytest.raises(config_service.AppTracingConfigInvalidConfigurationError): - ProviderConfigGateway(egress).prepare_new_config( + _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"}, @@ -165,65 +142,50 @@ def test_invalid_configs_are_rejected_before_any_request() -> None: def test_unknown_providers_are_rejected() -> None: - gateway = ProviderConfigGateway(RecordingEgress([], grpc_target=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() -> None: - payload = capabilities_payload("langfuse") - - assert payload["provider"] == "langfuse" - assert payload["native_cost"] is True - assert payload["export_ack"] == "accepted" - - -def test_empty_optional_secrets_are_stored_empty_and_not_reported_as_configured() -> None: - config = PhoenixConfig(api_key="", project="dify", endpoint="https://app.phoenix.arize.com") +def test_capabilities_describe_the_provider_mapping_and_the_collection_gate() -> None: + capabilities = TraceProviderCapabilitiesQuery( + providers=InstalledTraceProviders(LangfuseProvider), api_collection_enabled=True + ).get("langfuse") - stored = encode_config(WORKSPACE, config) - - assert stored["api_key"] == "" - assert configured_secrets(PhoenixProvider, stored) == () - assert decode_config(WORKSPACE, PhoenixProvider, stored) == config - - -def test_uninstalled_providers_can_still_be_named_but_not_presented(monkeypatch: pytest.MonkeyPatch) -> None: - registry.provider_class.cache_clear() - monkeypatch.setattr(registry, "_entry_points", dict) - gateway = ProviderConfigGateway(RecordingEgress([], grpc_target=None)) - try: - gateway.validate_provider("langfuse") - with pytest.raises(config_service.AppTracingConfigProviderUnavailableError): - gateway.present_config(workspace_id=WORKSPACE, tracing_provider="langfuse", tracing_config=_stored()) - finally: - registry.provider_class.cache_clear() - - -def test_secrets_that_cannot_be_decrypted_make_the_config_unreadable(monkeypatch: pytest.MonkeyPatch) -> None: - class _MissingWorkspaceKeyError(Exception): - pass - - def fail(tenant_id: str, tokens: list[str]) -> list[str]: - raise _MissingWorkspaceKeyError(f"{tenant_id}: {len(tokens)} tokens") - - monkeypatch.setattr(provider_configs, "batch_decrypt_token", fail) + 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 - with pytest.raises(ProviderConfigUnreadableError): - decode_config(WORKSPACE, LangfuseProvider, _stored()) +@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 + ) -def test_configs_without_secrets_never_load_the_workspace_key(monkeypatch: pytest.MonkeyPatch) -> None: - def forbidden(tenant_id: str, tokens: list[str]) -> list[str]: - raise AssertionError(f"{tenant_id} has no secret to decrypt, got {tokens}") + with pytest.raises(error): + query.get(provider_id) - monkeypatch.setattr(provider_configs, "batch_decrypt_token", forbidden) - config = decode_config(WORKSPACE, PhoenixProvider, {"api_key": None, "project": "dify"}) +def test_uninstalled_providers_can_still_be_named_but_not_presented() -> None: + gateway = ProviderConfigGateway( + providers=InstalledTraceProviders(), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + egress=RecordingEgress([], grpc_target=None), + ) - assert isinstance(config, PhoenixConfig) - assert config.api_key is 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( @@ -236,7 +198,7 @@ def test_stored_secrets_never_follow_an_edited_destination(submitted_secret: str submitted["secret_key"] = submitted_secret with pytest.raises(config_service.AppTracingConfigInvalidConfigurationError): - ProviderConfigGateway(egress).prepare_updated_config( + _gateway(egress).prepare_updated_config( workspace_id=WORKSPACE, tracing_provider="langfuse", tracing_config=submitted, @@ -248,7 +210,7 @@ def test_stored_secrets_never_follow_an_edited_destination(submitted_secret: str def test_an_edited_destination_accepts_re_entered_secrets() -> None: egress = RecordingEgress([_projects("project-2")], grpc_target=None) - stored = ProviderConfigGateway(egress).prepare_updated_config( + 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"}, @@ -262,7 +224,7 @@ def test_an_edited_destination_accepts_re_entered_secrets() -> None: def test_an_unchanged_destination_with_a_trailing_slash_keeps_secrets() -> None: egress = RecordingEgress([_projects("project-1")], grpc_target=None) - stored = ProviderConfigGateway(egress).prepare_updated_config( + stored = _gateway(egress).prepare_updated_config( workspace_id=WORKSPACE, tracing_provider="langfuse", tracing_config={"public_key": PUBLIC_KEY, "host": "https://langfuse.example.com/"}, 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 d3b3abf6415a80..301d70d8692181 100644 --- a/api/tests/unit_tests/services/tracing/test_recovery.py +++ b/api/tests/unit_tests/services/tracing/test_recovery.py @@ -1,20 +1,25 @@ -"""Recover lost terminal notifications from real business rows through the trace assembler.""" +"""Recover lost terminal notifications from real business rows through the trace record repository.""" import json from collections.abc import Callable, Generator, Iterator -from dataclasses import dataclass -from datetime import datetime, timedelta +from dataclasses import dataclass, replace +from datetime import UTC, datetime, timedelta from decimal import Decimal from typing import Literal, cast, override +from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 import pytest 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 +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 @@ -46,22 +51,36 @@ from core.entities.provider_entities import ProviderQuotaType, QuotaConfiguration, QuotaUnit, SystemConfiguration from events.event_handlers import update_provider_when_message_created as receiver from extensions.ext_database import db -from extensions.tracing import bind_tracing +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.account import Account, TenantAccountJoin +from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.application_trace import ApplicationTraceInvocation, ApplicationTraceRun -from models.enums import CreatorUserRole, MessageStatus, WorkflowRunTriggeredFrom -from models.model import App, AppMode, Conversation, Message +from models.enums import ConversationFromSource, CreatorUserRole, MessageStatus, WorkflowRunTriggeredFrom +from models.model import App, AppMode, Conversation, Message, TraceAppConfig from models.workflow import WorkflowRun, WorkflowType -from repositories.application_trace_repository import TraceCollectionError -from services.tracing import recovery, runtime_bridge -from services.tracing.assembly import SourceNotReadyError -from services.tracing.emitter import TraceScopes +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 StoredPolicyTraceScopes +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import AgentTraceSettings, TraceInvocation, TraceRun, TraceRunState +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.config_override import apply_config_overrides 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, @@ -72,35 +91,232 @@ _StoppedQueueManager, _UsageCancellationClient, ) -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_accepted as accepted -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_install_delivery as install_delivery, -) -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_tracing_service as tracing_service, -) -from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection -from tests.unit_tests.services.tracing.test_service import ( - add_app_owner, - attach_conversation, - message_record, - sent_traces, -) -@dataclass +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: + return self.repository.invocation( + run_id=self.run.id, 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 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 = repository.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 message is not None + session.add( + make_conversation( + conversation_id=message.conversation_id, + 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) -> tuple[int, str]: + return collection.repository.ingest(batch, max_spans=100, max_bytes=1_000_000) + + +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_terminal( + run, + traces=self.collection.repository, + records=self.records, + collector=self.service, + collection_timeout=self.settings.collection_timeout_seconds, + 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 = TraceRecordRepository(collection.sessions, files_url="https://files.test", enterprise_details=False) + scheduled: list[str] = [] + service = ApplicationTracingService( + traces=collection.repository, + records=records, + builder=CanonicalTraceBuilder(), + providers=InstalledTraceProviders(ScriptedProvider), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + egress=destination, + schedule_export=scheduled.append, + settings=settings, + ) + return RecoveryCase(collection, records, service, settings, now, scheduled, destination) + class _Recorder: def __init__(self) -> None: @@ -125,23 +341,32 @@ def _publish(self, event: AppQueueEvent, pub_from: PublishFrom) -> None: raise AssertionError("Stream dispatch must not publish queue events") -def recording_service( - collection: Collection, destination: RecordingEgress, scheduled: list[str], monkeypatch: pytest.MonkeyPatch -) -> ApplicationTracingService: - """The ledger the runtime hooks reach; scheduled exports are recorded instead of enqueued.""" - service = tracing_service(collection, destination) - monkeypatch.setattr(service, "schedule", scheduled.append) - monkeypatch.setattr(runtime_bridge, "_service", lambda: service) - return service +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] = [] -def tracing_application(collection: Collection) -> Flask: - """An application bound like the composition root, with events kept in memory.""" + @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=TraceScopes(emitter=_Recorder(), session_factory=collection.sessions, enterprise_enabled=lambda: False), - hooks=runtime_bridge.RuntimeApplicationTracingHooks(), + scopes=StoredPolicyTraceScopes( + TraceScopes( + emitter=_Recorder(), + policies=SQLAlchemyAppTracingConfigRepository(session_factory=case.collection.sessions), + enterprise_enabled=lambda: False, + ) + ), + hooks=RuntimeApplicationTracingHooks(case.service), ) return application @@ -150,10 +375,10 @@ def tracing_application(collection: Collection) -> Flask: def case( sqlite_engine: Engine, sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch ) -> RecoveryCase: - collection = seed(sqlite_session_factory, "recover", capture_content=True) + collection = seed(sqlite_session_factory, "recover", capture_content=True, capture_agent=True) attach_conversation(collection, AppMode.AGENT) - invocation = collection.invocation() - ingest(collection, batch_for(collection.run, invocation)) + 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) @@ -169,12 +394,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(recovery, "naive_utc_now", lambda: now) - install_delivery(monkeypatch) - destination = accepted("provider-id") - scheduled: list[str] = [] - service = recording_service(collection, destination, scheduled, monkeypatch) - return RecoveryCase(collection, service, now, scheduled, destination) + 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: @@ -199,7 +455,6 @@ def save_agent_stop( notify: Callable[[Message, ChatAppGenerateEntity], None] | None = None, ) -> list[TraceEvent]: """Run the real consumer save against the owned SQL rows; returns the trace events it emitted.""" - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) def send(message: Message, *, application_generate_entity: ChatAppGenerateEntity) -> None: if notify is not None: @@ -295,8 +550,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, @@ -324,41 +579,40 @@ def notify(message: Message, entity: ChatAppGenerateEntity) -> None: 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: destination = RecordingEgress( [ response(200, {"trace_id": "provider-id"}) - if delivery == "accepted" + if delivery is TraceRunState.ACCEPTED else EgressError("acknowledgement lost", request_sent=True) ], grpc_target=None, ) - service = recording_service(case.collection, destination, case.scheduled, monkeypatch) + 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 = message_record(service, case.collection) + notification = message_record(stopped) queued = save_agent_stop(case, monkeypatch, answer) - assert service.capture_terminal(notification) 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(service, stored) is False + assert stopped.recover(stored, now=case.now) is False assert case.collection.repository.claim(stored.id, collection_timeout=60, lease_seconds=30) is None service.deliver(stored.id) assert destination.requests == [] @@ -371,21 +625,20 @@ def test_stopped_consumer_cannot_freeze_a_terminal_before_the_runtime_worker( outcome=outcome, output=output, ) - terminal = case.collection.repository.get(stored.id) - assert terminal is not None - trace = service.assemble(terminal) - assert trace.root.attributes["message_execution_state"] == outcome - assert trace.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. service.deliver(stored.id) service.deliver(stored.id) - assert sent_traces(destination) == [trace] + exported = sent_traces(destination) + assert len(exported) == 1 + 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 + assert frozen.state is delivery service.capture_agent_app_terminal( tenant_id=stored.tenant_id, app_id=stored.app_id, @@ -393,7 +646,7 @@ def test_stopped_consumer_cannot_freeze_a_terminal_before_the_runtime_worker( outcome="succeeded" if outcome != "succeeded" else "cancelled", output="late result must not replace the accepted envelope", ) - assert service.capture_terminal(notification) 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 @@ -436,10 +689,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 @@ -463,10 +716,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") @@ -484,7 +737,7 @@ def test_stopped_runner_hands_off_only_an_observed_runtime_terminal( snapshot_save_fails: bool, ) -> None: 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") @@ -527,7 +780,7 @@ def unexpected_cancel(*_args: object, **_kwargs: object) -> None: agent_config_snapshot_id=identifier("stop-snapshot"), home_snapshot_id=None, ) - with tracing_application(case.collection).app_context(), pytest.raises(GenerateTaskStoppedError): + with tracing_application(case).app_context(), pytest.raises(GenerateTaskStoppedError): runner._consume_stream( client.run_id, dify_context=context, @@ -543,19 +796,18 @@ 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) - assert trace.root.attributes["message_execution_state"] == runtime_outcome - assert trace.root.status.value == ("ok" if runtime_outcome == "succeeded" else "error") + 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) if pending_invocation is not None: assert case.collection.repository.claim(run.id, collection_timeout=60, lease_seconds=30) is None @@ -566,32 +818,32 @@ def unexpected_cancel(*_args: object, **_kwargs: object) -> None: 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_run = case.collection.repository.get(run.id) assert accepted_run is not None - assert accepted_run.state == "accepted" + 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: - 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" ) @@ -644,7 +896,7 @@ def publish_after_stop(event: AppQueueEvent, pub_from: PublishFrom) -> None: user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, ) - with tracing_application(case.collection).app_context(), pytest.raises(GenerateTaskStoppedError): + with tracing_application(case).app_context(), pytest.raises(GenerateTaskStoppedError): runner.run( dify_context=context, agent_id=identifier("stop-agent"), @@ -677,7 +929,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: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) client = _UsageCancellationClient() queue_manager = _FakeQueueManager() @@ -712,7 +963,7 @@ def cancel_then_stop(*_args: object, **_kwargs: object) -> RunEvent: agent_config_snapshot_id=identifier("cleanup-snapshot"), home_snapshot_id=None, ) - with tracing_application(case.collection).app_context(): + with tracing_application(case).app_context(): with pytest.raises(GenerateTaskStoppedError if stop_during_cleanup else RuntimeError): runner._consume_stream( client.run_id, @@ -727,7 +978,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, @@ -751,7 +1004,7 @@ def cancel_then_stop(*_args: object, **_kwargs: object) -> RunEvent: case.service.deliver(run.id) delivered = case.collection.repository.get(run.id) assert delivered is not None - assert delivered.state == "accepted" + assert delivered.state is TraceRunState.ACCEPTED assert delivered.attempts == 1 assert len(case.destination.requests) == 1 @@ -760,24 +1013,24 @@ def cancel_then_stop(*_args: object, **_kwargs: object) -> RunEvent: def test_missing_stop_worker_without_runtime_spans_also_reaches_a_delivery_failure( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, capture_agent: bool ) -> None: - collection = seed(case.collection.sessions, f"stop-without-spans-{capture_agent}", capture_agent=capture_agent) + collection = seed( + case.collection.sessions, + f"stop-without-spans-{capture_agent}", + capture_content=False, + capture_agent=capture_agent, + ) attach_conversation(collection, AppMode.AGENT) - scheduled: list[str] = [] - destination = RecordingEgress([], grpc_target=None) - stopped = RecoveryCase( + stopped = recovery_case( collection, - recording_service(collection, destination, scheduled, monkeypatch), + RecordingEgress([], grpc_target=None), + replace(AGENT_TRACE_SETTINGS, collection_timeout_seconds=60), case.now, - scheduled, - destination, ) - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_COLLECTION_TIMEOUT=60) 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 @@ -792,7 +1045,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, @@ -817,7 +1070,7 @@ def test_worker_terminal_survives_closed_sse_without_inventing_a_message_write( def test_worker_terminal_revalidates_message_tenant(case: RecoveryCase) -> None: - 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, @@ -830,19 +1083,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_the_assembler(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) @@ -865,7 +1118,7 @@ def test_deliver_recovers_and_exports_once_after_terminal_queue_loss(case: Recov 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 @@ -877,24 +1130,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 @@ -905,13 +1158,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 @@ -929,7 +1182,7 @@ 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 @@ -941,58 +1194,61 @@ def test_explicit_message_failure_can_recover_without_fabricating_usage(case: Re assert [(span["usage"], span["cost"]) for span in stored.outer_trace["spans"]] == [(None, None)] -def test_cross_tenant_source_is_rejected_before_the_assembler_runs( - 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[TraceEvent] = [] - - def forbidden(event: TraceEvent) -> None: - called.append(event) - - monkeypatch.setattr(case.service.assembler, "assemble", forbidden) - assert recovery.recover_terminal(case.service, case.collection.run) is False - assert called == [] + records = _FailingRecords(AssertionError("a foreign source must never be assembled")) + recovered = recover_terminal( + case.collection.run, + traces=case.collection.repository, + records=records, + collector=case.service, + collection_timeout=case.settings.collection_timeout_seconds, + 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 test_source_not_ready_for_assembly_is_retried_later(case: RecoveryCase, monkeypatch: pytest.MonkeyPatch) -> None: +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) - events: list[TraceEvent] = [] - - def not_ready(event: TraceEvent) -> None: - events.append(event) - raise SourceNotReadyError - - monkeypatch.setattr(case.service.assembler, "assemble", not_ready) - assert recovery.recover_terminal(case.service, case.collection.run) is False - assert [event.payload for event in events] == [ + records = _FailingRecords(SourceNotReadyError()) + recovered = recover_terminal( + case.collection.run, + traces=case.collection.repository, + records=records, + collector=case.service, + collection_timeout=case.settings.collection_timeout_seconds, + 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 events[0].correlation.session_id == "session-original" + 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.state, stored.error_code, stored.outer_trace) == ("collecting", None, None) + 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]) @@ -1002,7 +1258,7 @@ 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 @@ -1036,13 +1292,14 @@ def add_workflow_run(case: RecoveryCase, status: WorkflowExecutionStatus, *, fin return workflow_run_id -def start_workflow(case: RecoveryCase, workflow_run_id: str) -> ApplicationTraceRun: +def start_workflow(case: RecoveryCase, workflow_run_id: str) -> TraceRun: 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", + workflow_id=None, ) assert run is not None return run @@ -1072,7 +1329,7 @@ def test_workflow_recovery_requires_committed_terminal_state( 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: @@ -1087,7 +1344,7 @@ def test_workflow_recovery_requires_committed_terminal_state( assert stored.outer_trace["user_id"] == owner_id else: assert stored.outer_trace is None - assert (stored.state, stored.error_code) == ("collecting", None) + assert (stored.state, stored.error_code) == (TraceRunState.COLLECTING, None) @pytest.mark.parametrize("owner_state", ["no_creator", "missing_account", "not_a_member"]) @@ -1106,7 +1363,7 @@ def test_workflow_recovery_does_not_depend_on_the_app_creator(case: RecoveryCase 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 recovery.recover_terminal(case.service, run) is True + 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 @@ -1143,7 +1400,8 @@ def test_early_error_root_preserves_known_query_only_when_content_is_enabled( 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.span_id != assembled.root.span_id) with case.collection.sessions() as session: @@ -1229,7 +1487,7 @@ def test_worker_handoff_leaves_a_not_collected_run_untouched(case: RecoveryCase) with case.collection.sessions.begin() as session: run = session.get(ApplicationTraceRun, case.collection.run.id) assert run is not None - run.state = "not_collected" + 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, @@ -1241,4 +1499,4 @@ def test_worker_handoff_leaves_a_not_collected_run_untouched(case: RecoveryCase) assert stored is not None assert stored.outer_trace is None assert case.scheduled == [] - assert claim(case.collection) is None + assert case.collection.repository.claim(case.collection.run.id, collection_timeout=60, lease_seconds=30) is None diff --git a/api/tests/unit_tests/services/tracing/test_registry.py b/api/tests/unit_tests/services/tracing/test_registry.py index 71f36e1b1e2ebc..c02662bbd16db8 100644 --- a/api/tests/unit_tests/services/tracing/test_registry.py +++ b/api/tests/unit_tests/services/tracing/test_registry.py @@ -4,32 +4,28 @@ import pytest from services.tracing import registry -from services.tracing.registry import ( - ENTRY_POINT_GROUP, - SUPPORTED_PROVIDERS, - TraceProviderUnavailableError, - UnknownTraceProviderError, - installed_providers, - provider_class, -) +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]: - provider_class.cache_clear() + registry._load.cache_clear() yield - provider_class.cache_clear() + registry._load.cache_clear() def test_every_supported_provider_is_installed_under_its_own_id() -> None: - assert installed_providers() == SUPPORTED_PROVIDERS + catalog = EntryPointTraceProviderCatalog() + for provider_id in SUPPORTED_PROVIDERS: - assert provider_class(provider_id).id == provider_id + assert catalog.resolve(provider_id).id == provider_id def test_unknown_providers_are_client_errors() -> None: with pytest.raises(UnknownTraceProviderError): - provider_class("datadog") + EntryPointTraceProviderCatalog().resolve("datadog") @pytest.mark.parametrize( @@ -52,5 +48,4 @@ def test_supported_providers_that_do_not_load_are_unavailable( monkeypatch.setattr(registry, "_entry_points", lambda: entry_points) with pytest.raises(TraceProviderUnavailableError): - provider_class("langfuse") - assert "langfuse" not in installed_providers() + EntryPointTraceProviderCatalog().resolve("langfuse") 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 ed931da6c63a42..72853cef81ef10 100644 --- a/api/tests/unit_tests/services/tracing/test_runtime_bridge.py +++ b/api/tests/unit_tests/services/tracing/test_runtime_bridge.py @@ -1,13 +1,17 @@ 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 Literal, 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 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, @@ -17,9 +21,11 @@ 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 sqlalchemy import delete, func, select -from sqlalchemy.exc import OperationalError +from flask import Flask +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 @@ -27,8 +33,10 @@ from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom, WorkflowAppGenerateEntity from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeBuildContext +from extensions.tracing import application_tracing_hooks, bind_tracing from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus from graphon.runtime import VariablePool +from models.account import Tenant, TenantAccountJoin, TenantAccountRole from models.agent import ( Agent, AgentConfigSnapshot, @@ -39,8 +47,7 @@ ) from models.agent_config_entities import AgentSoulConfig, WorkflowNodeJobConfig from models.application_trace import ApplicationTraceRun -from models.enums import 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, @@ -48,36 +55,41 @@ 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.application_trace_repository import ApplicationTraceRepository, TraceCollectionError -from services.tracing import runtime_bridge -from services.tracing.assembly import TraceAssembler +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 StoredPolicyTraceScopes from services.tracing.canonical import CanonicalTraceBuilder -from services.tracing.delivery import TraceCollectorUnavailableError, TraceDeliveryService +from services.tracing.collection import ( + AgentTraceSettings, + ApplicationTraceStore, + TraceRun, + TraceRunState, + TraceSourceKind, + workflow_has_agent_v2, +) +from services.tracing.delivery import ExportJob, ExportRetryQueue, TraceDeliveryService +from services.tracing.errors import TraceCollectorUnavailableError, TracingStoreUnavailableError from services.tracing.normalization import remap_outer_trace +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 tests.helpers.unit.config_override import apply_config_overrides -from tests.helpers.unit.model_factories import make_account, make_app, 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_no_egress as no_egress -from tests.unit_tests.services.tracing.test_service import ( - _application_trace_collection_seed as _application_trace_handoff_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 ( - _application_trace_collection_tracing_service as tracing_service, +from tests.helpers.unit.model_factories import ( + make_account, + make_app, + make_conversation, + make_message, + make_tenant, + make_workflow, ) -from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection -from tests.unit_tests.services.tracing.test_service import ( - add_app_owner, - attach_conversation, - install_decryption, - trace_event, +from tests.helpers.unit.services.tracing_ports import ( + AGENT_TRACE_SETTINGS, + InstalledTraceProviders, + WorkspaceBoundCipher, ) """SQLite ownership regressions at the handoff from trace delivery to Agent trace collection.""" @@ -92,19 +104,272 @@ STARTED_AT = datetime(2026, 10, 15, tzinfo=UTC) -@dataclass +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 freezing a run fails with ``failure``.""" + + def __init__(self, sessions: sessionmaker[Session], failure: Exception) -> None: + super().__init__(sessions) + self._failure = failure + + @override + def start( + self, + *, + tenant_id: str, + app_id: str, + source_kind: TraceSourceKind, + source_id: str, + session_id: str, + workflow_id: str | None, + ) -> TraceRun | None: + raise self._failure + + +def record_source(sessions: sessionmaker[Session]) -> TraceRecordRepository: + return TraceRecordRepository(sessions, files_url="https://files.test", enterprise_details=False) + + +def tracing_service( + sessions: sessionmaker[Session], + traces: ApplicationTraceStore, + egress: Egress, + scheduled: list[str], + settings: AgentTraceSettings, +) -> ApplicationTracingService: + """Agent trace collection exporting to Langfuse; scheduled exports are recorded instead of enqueued.""" + return ApplicationTracingService( + traces=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 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: + 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 = repository.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 message is not None + session.add( + make_conversation( + conversation_id=message.conversation_id, + 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.""" + 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 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,), + ) + collection.repository.terminal(run.id, trace.model_dump(mode="json"), root.span_id) + + +@dataclass(frozen=True) class Case: collection: Collection service: ApplicationTracingService scheduled: list[str] -def setup( - sessions: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, name: str, mode: AppMode, *, snapshot: bool -) -> Case: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - install_decryption(monkeypatch) - collection = _application_trace_handoff_seed(sessions, name) +def handoff(collection: Collection, traces: ApplicationTraceStore, egress: Egress) -> Case: + """The Agent collector over ``traces``, exporting through ``egress``.""" + scheduled: list[str] = [] + return Case( + collection, tracing_service(collection.sessions, traces, egress, scheduled, AGENT_TRACE_SETTINGS), scheduled + ) + + +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) @@ -112,43 +377,26 @@ def setup( app.mode = mode if not snapshot: session.execute(delete(ApplicationTraceRun).where(ApplicationTraceRun.id == collection.run.id)) - service = tracing_service(collection, no_egress()) - scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) - monkeypatch.setattr(runtime_bridge, "_service", lambda: service) - return Case(collection, service, scheduled) + 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 ledger as the Agent collector.""" + """Deliver one event the way the trace worker does, with this case's service as the Agent collector.""" + sessions = case.collection.sessions TraceDeliveryService( - session_factory=case.collection.sessions, - assembler=case.service.assembler, + 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 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 message(case: Case) -> Message: with case.collection.sessions() as session: row = session.get(Message, case.collection.run.source_id) @@ -158,7 +406,7 @@ def message(case: Case) -> Message: def message_event(case: Case) -> TraceEvent: row = message(case) - return trace_event(case.collection, MessageCompleted(message_id=row.id, conversation_id=row.conversation_id)) + return trace_event(case.collection.run, MessageCompleted(message_id=row.id, conversation_id=row.conversation_id)) def workflow_event(row: WorkflowRun, parent: WorkflowRun | None) -> TraceEvent: @@ -178,9 +426,7 @@ def workflow_event(row: WorkflowRun, parent: WorkflowRun | None) -> TraceEvent: ) -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: @@ -199,13 +445,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), @@ -222,6 +468,7 @@ def workflow( source_kind="workflow", source_id=row.id, session_id="workflow-session", + workflow_id=None, ) if snapshot else None @@ -230,6 +477,7 @@ def workflow( 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"] = { @@ -237,7 +485,6 @@ def workflow_info(row: WorkflowRun, parent: WorkflowRun | None) -> WorkflowTrace "parent_node_execution_id": identifier("parent-tool-node"), } return WorkflowTraceInfo( - workflow_data=row, node_executions=[], conversation_id=None, message_started_at=None, @@ -263,17 +510,10 @@ def workflow_info(row: WorkflowRun, parent: WorkflowRun | None) -> WorkflowTrace ) -def run_count(case: Case) -> int | None: - with case.collection.sessions() as session: - return session.scalar(select(func.count()).select_from(ApplicationTraceRun)) - - class TestApplicationTraceHandoffCases: @staticmethod - def test_chatflow_message_event_is_not_exported_again( - 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)) @@ -290,19 +530,17 @@ def test_chatflow_message_event_is_not_exported_again( assert after is not None assert after.outer_trace == before.outer_trace assert case.scheduled == [] - assert run_count(case) == 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, @@ -320,9 +558,9 @@ def fail_start(**_kwargs: object) -> None: binding_id=identifier("binding"), backend_binding_ref="backend", ) - assert runtime_bridge.prepare_agent_app(context) is None + assert RuntimeApplicationTracingHooks(failing.service).prepare_agent_app(context) is None destination = no_egress() - deliver(case, message_event(case), destination) + deliver(failing, message_event(failing), destination) assert destination.requests == [] stored = case.collection.repository.source_run( tenant_id=case.collection.run.tenant_id, @@ -331,7 +569,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 == {} @@ -339,9 +577,9 @@ 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) + 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. @@ -383,25 +621,25 @@ def test_auxiliary_events_resolve_owner_without_bypassing_content_policy( ) ) destination = exported() if mode is AppMode.CHAT else no_egress() - deliver(case, trace_event(case.collection, payload), destination) + 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) == 0 + assert run_count(case.collection.sessions) == 0 @staticmethod def test_plain_nested_workflow_keeps_its_parent_context_for_the_provider( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + 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) add_app_owner(child_case.collection, "plain-child") event = workflow_event(child, parent) - record = child_case.service.assembler.assemble(event) + record = record_source(sqlite_session_factory).assemble(event) assert isinstance(record, WorkflowTraceInfo) - assert child_case.service.capture_terminal(record) is False - assert child_case.service.link_managed_parent(record) is record + 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") @@ -410,13 +648,13 @@ def test_plain_nested_workflow_keeps_its_parent_context_for_the_provider( 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(child_case) == 0 + assert run_count(sqlite_session_factory) == 0 @staticmethod def test_plain_chatflow_auxiliary_events_keep_the_provider_route( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + 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: stored = session.get(Message, case.collection.run.source_id) @@ -431,21 +669,21 @@ def test_plain_chatflow_auxiliary_events_keep_the_provider_route( ended_at=STARTED_AT, ) destination = exported() - deliver(case, trace_event(case.collection, payload), destination) + deliver(case, trace_event(case.collection.run, payload), destination) assert len(destination.requests) == 1 @staticmethod def test_provider_exported_child_of_managed_parent_names_the_collected_parent( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + 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, "provider-child", AppMode.WORKFLOW, 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 - linked = child_case.service.link_managed_parent(info) + 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") @@ -464,24 +702,27 @@ def test_provider_exported_child_of_managed_parent_names_the_collected_parent( @staticmethod def test_collector_database_failure_redelivers_without_provider_fallback( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + sqlite_session_factory: sessionmaker[Session], tmp_path: Path ) -> None: - case = setup(sqlite_session_factory, monkeypatch, "collector-unavailable", AppMode.AGENT, snapshot=True) - - def unavailable(**_kwargs: object) -> None: - raise OperationalError("statement", {}, Exception("sensitive database error")) - - monkeypatch.setattr(case.collection.repository, "source_run", unavailable) + 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() - with pytest.raises(TraceCollectorUnavailableError): - deliver(case, message_event(case), destination) + 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") @@ -573,17 +814,15 @@ def test_mixed_workflow_exports_normal_llm_and_agent_calls_without_double_usage( max_spans=100, max_bytes=100000, ) - record = case.service.assembler.assemble(workflow_event(row, None)) + record = record_source(sqlite_session_factory).assemble(workflow_event(row, None)) assert isinstance(record, WorkflowTraceInfo) - assert case.service.capture_terminal(record) is True + assert case.service.capture_terminal(tenant_id=row.tenant_id, info=record) is True assert case.scheduled == [managed.id] destination = exported() - tracing_service(Collection(case.collection.repository, case.collection.sessions, managed), destination).deliver( - managed.id - ) + 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 == "accepted" + 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)] @@ -595,21 +834,23 @@ def test_mixed_workflow_exports_normal_llm_and_agent_calls_without_double_usage( """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: ApplicationTraceStore, settings: AgentTraceSettings +) -> RuntimeApplicationTracingHooks: + return RuntimeApplicationTracingHooks(tracing_service(sessions, traces, no_egress(), [], settings)) def run_context(app_id: str) -> DifyRunContext: @@ -636,7 +877,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( @@ -671,215 +912,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), - assembler=TraceAssembler(sqlite_session_factory, files_url="https://files.test", enterprise_details=False), - builder=CanonicalTraceBuilder(), - egress=no_egress(), + 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) ) - 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 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_agent_app_terminal", agent_terminal) - monkeypatch.setattr(runtime_bridge, "_service", lambda: target) - return case @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, ) - runtime_bridge.capture_agent_app_terminal( - tenant_id="tenant", app_id="app", message_id="message", outcome="succeeded", output="answer" - ) - 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_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: - runtime_bridge.capture_agent_app_terminal( + def test_worker_terminal_handoff_preserves_outcome_without_sse( + bridge: Bridge, outcome: Literal["succeeded", "failed", "cancelled"] + ) -> None: + started = bridge.traces.start( + tenant_id="tenant", + app_id="agent-app", + 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", + 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: - 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", @@ -891,34 +1142,27 @@ def test_workflow_fail_open_logs_error_kind_without_content( assert "private-prompt" not in caplog.text @staticmethod - def test_hooks_expose_only_the_agent_collection_lifecycle() -> None: - hooks = runtime_bridge.RuntimeApplicationTracingHooks() - assert hooks.prepare_agent_app is runtime_bridge.prepare_agent_app - assert hooks.prepare_workflow_agent is runtime_bridge.prepare_workflow_agent - assert hooks.freeze_workflow_run is runtime_bridge.freeze_workflow_run - assert hooks.capture_agent_app_terminal is runtime_bridge.capture_agent_app_terminal - assert {name for name in dir(hooks) if not name.startswith("_")} == { + 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 [ @@ -946,6 +1190,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), @@ -967,20 +1219,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="{}", @@ -1024,26 +1270,20 @@ def _workflow_tool_tracing_boundary_seed( ) node_record.created_at = now session.add(node_record) - service = ApplicationTracingService( - ApplicationTraceRepository(sessions), - assembler=TraceAssembler(sessions, files_url="https://files.test", enterprise_details=False), - builder=CanonicalTraceBuilder(), - egress=no_egress(), - ) - 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( @@ -1053,7 +1293,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: @@ -1061,38 +1301,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] = [] + statements: list[str] = [] + + def executed( + _connection: object, + _cursor: object, + statement: str, + _parameters: object, + _context: object, + _executemany: object, + ) -> None: + statements.append(statement) - def lookup() -> ApplicationTracingService: - service_lookups.append(True) - return service - - 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) @@ -1102,7 +1351,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 @@ -1111,12 +1360,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( @@ -1125,9 +1373,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) @@ -1137,6 +1385,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, @@ -1146,7 +1395,6 @@ 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, @@ -1173,14 +1421,40 @@ def test_hidden_child_skip_keeps_normal_root_agent_trace_assemblable( outer = CanonicalTraceBuilder().build(info) assert outer is not None outer = remap_outer_trace(outer, trace_id=run.trace_id, capture_content=False, synthetic_span_ids=frozenset()) - 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) + traces.terminal(run.id, outer.model_dump(mode="json"), outer.root.span_id) + [invocation] = traces.invocations(run.id) + traces.ingest( + 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.span_id for span in trace.spans} - assert all(span.parent_span_id is None or span.parent_span_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 0354bff549d33a..d113757b5a270a 100644 --- a/api/tests/unit_tests/services/tracing/test_service.py +++ b/api/tests/unit_tests/services/tracing/test_service.py @@ -1,10 +1,10 @@ import json import threading -from collections.abc import Callable, Iterator, Mapping -from dataclasses import dataclass +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass, field, replace from datetime import UTC, datetime, timedelta from decimal import Decimal -from typing import override +from typing import NoReturn, override from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 import pytest @@ -20,7 +20,9 @@ from dify_trace_sdk.testing import RecordingEgress, response from sqlalchemy import Connection, Engine, event, func, select from sqlalchemy.orm import Session, sessionmaker +from typing_extensions import TypedDict +from core.helper.encrypter import obfuscated_token from extensions.logstore.aliyun_logstore import AliyunLogStore from extensions.logstore.repositories.logstore_api_workflow_node_execution_repository import ( _dict_to_workflow_node_execution_model, @@ -30,21 +32,34 @@ from libs.datetime_utils import naive_utc_now from machinery.context import RequestContext from models.account import Tenant, TenantAccountJoin, TenantAccountRole -from models.application_trace import ApplicationTraceBatch, ApplicationTraceInvocation, ApplicationTraceRun +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, WorkflowRun from repositories import application_trace_repository as repository_module from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError +from repositories.application_trace_repository import ApplicationTraceRepository +from repositories.trace_record_repository import TraceRecordRepository from services.app_tracing_config_service import AppTracingConfigService -from services.tracing import provider_configs as provider_configs_module -from services.tracing import service as service_module -from services.tracing.assembly import TraceAssembler from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import ( + AgentTraceSettings, + TraceInvocation, + TraceRun, + TraceRunState, + TraceRunStatus, + TraceSourceKind, +) +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 UnknownTraceProviderError, provider_class +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 ( @@ -55,6 +70,7 @@ 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.""" @@ -67,9 +83,9 @@ def _application_trace_collection_identifier(value: str) -> str: class _ApplicationTraceCollectionCollection: repository: ApplicationTraceRepository sessions: sessionmaker[Session] - run: ApplicationTraceRun + run: TraceRun - def invocation(self, name: str = "first") -> ApplicationTraceInvocation: + def invocation(self, name: str = "first") -> TraceInvocation: return self.repository.invocation( run_id=self.run.id, agent_id=_application_trace_collection_identifier("agent"), @@ -112,28 +128,89 @@ def send(self, config: LangfuseConfig, trace: Trace) -> SendResult: ) -def install_decryption(monkeypatch: pytest.MonkeyPatch) -> list[str]: - """A stored ``encrypted-x`` secret decrypts to ``plaintext-x``; returns the ciphertexts decrypted.""" - ciphertexts: list[str] = [] +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 decrypt(_tenant_id: str, tokens: list[str]) -> list[str]: - ciphertexts.extend(tokens) + 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] - monkeypatch.setattr(provider_configs_module, "batch_decrypt_token", decrypt) - return ciphertexts + @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) -def _application_trace_collection_install_delivery(monkeypatch: pytest.MonkeyPatch) -> list[str]: - """Deliver Agent traces through ``ScriptedProvider``; returns the ciphertexts decrypted.""" +class UnreadableSecretCipher(StoredSecretCipher): + """No stored secret can be decrypted, as when the workspace key is lost.""" - def provider_class(provider_id: str) -> type[ScriptedProvider]: - if provider_id != ScriptedProvider.id: - raise UnknownTraceProviderError(provider_id) - return ScriptedProvider + @override + def decrypt(self, tenant_id: str, tokens: Sequence[str]) -> NoReturn: + raise ValueError("private-ciphertext") - monkeypatch.setattr(service_module, "provider_class", provider_class) - return install_decryption(monkeypatch) + +def trace_records(sessions: sessionmaker[Session]) -> TraceRecordRepository: + return TraceRecordRepository(sessions, files_url="https://files.test", enterprise_details=False) + + +@dataclass +class ServicePorts: + """The service's collaborators besides its ledger; 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 ApplicationTracingService( + traces=case.repository, + records=trace_records(case.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: @@ -150,17 +227,6 @@ 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_tracing_service( - case: _ApplicationTraceCollectionCollection, egress: Egress -) -> ApplicationTracingService: - return ApplicationTracingService( - case.repository, - assembler=TraceAssembler(case.sessions, files_url="https://files.test", enterprise_details=False), - builder=CanonicalTraceBuilder(), - egress=egress, - ) - - def _application_trace_collection_seed( sessions: sessionmaker[Session], name: str, @@ -219,7 +285,12 @@ def _application_trace_collection_seed( ) 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" + 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 _ApplicationTraceCollectionCollection(repository, sessions, run) @@ -269,13 +340,13 @@ def trace_event(case: _ApplicationTraceCollectionCollection, payload: TracePaylo ) -def message_record(service: ApplicationTracingService, case: _ApplicationTraceCollectionCollection) -> MessageTraceInfo: +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 = service.assembler.assemble( + record = trace_records(case.sessions).assemble( trace_event(case, MessageCompleted(message_id=case.run.source_id, conversation_id=conversation_id)) ) assert isinstance(record, MessageTraceInfo) @@ -283,8 +354,8 @@ def message_record(service: ApplicationTracingService, case: _ApplicationTraceCo def _application_trace_collection_batch_for( - run: ApplicationTraceRun, - invocation: ApplicationTraceInvocation, + run: TraceRun, + invocation: TraceInvocation, *, final: bool = True, batch_id: UUID | None = None, @@ -340,7 +411,7 @@ def _application_trace_collection_terminal(collection: _ApplicationTraceCollecti def _application_trace_collection_claim( collection: _ApplicationTraceCollectionCollection, -) -> tuple[ApplicationTraceRun, str] | None: +) -> tuple[TraceRun, str] | None: return collection.repository.claim(collection.run.id, collection_timeout=60, lease_seconds=30) @@ -364,8 +435,10 @@ def test_batch_duplicate_is_idempotent_after_final_and_acceptance( collection.repository.complete( collection.run.id, leased[1], - state="accepted", + state=TraceRunState.ACCEPTED, + error_code=None, result=SendResult(trace_id="provider-id", trace_url=None, span_urls={}), + retry_delay=None, ) assert _application_trace_collection_ingest(collection, batch) == (1, collection.run.id) stored = collection.repository.get(collection.run.id) @@ -381,18 +454,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.collected(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 @@ -400,7 +474,7 @@ 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"): + with pytest.raises(TraceCollectionLimitError, 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 @@ -413,7 +487,7 @@ 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 @@ -454,7 +528,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 @@ -464,7 +538,7 @@ 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"): + with pytest.raises(TraceCollectionConflictError, match="tracing_invocation_conflict"): collection.repository.invocation( run_id=collection.run.id, agent_id=_application_trace_collection_identifier("different-agent"), @@ -507,19 +581,21 @@ 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" + assert current.state is TraceRunState.UNCERTAIN collection.repository.complete( collection.run.id, leased[1], - state="accepted", + state=TraceRunState.ACCEPTED, + error_code=None, result=SendResult(trace_id="late", trace_url=None, span_urls={}), + retry_delay=None, ) - 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" + 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 @@ -527,13 +603,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 @@ -554,6 +630,7 @@ def test_running_trace_freezes_credentials_content_policy_and_session( source_kind="message", source_id=original.source_id, session_id="changed", + workflow_id=None, ) assert repeated is not None assert repeated.provider == "langfuse" @@ -566,34 +643,36 @@ def test_running_trace_freezes_credentials_content_policy_and_session( @pytest.mark.parametrize( ("outcome", "reply"), [ - ("accepted", response(200, {"trace_id": "vendor-trace", "trace_url": "https://trace.test/example"})), - ("uncertain", EgressError("acknowledgement lost", request_sent=True)), - ("retrying", response(503, {})), + ( + 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, + 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) - ciphertexts = _application_trace_collection_install_delivery(monkeypatch) + cipher = StoredSecretCipher() egress = RecordingEgress([reply], grpc_target=None) - service = _application_trace_collection_tracing_service(collection, egress) + service = ServicePorts(egress, cipher=cipher).service(collection) service.deliver(collection.run.id) service.deliver(collection.run.id) 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 ciphertexts == ["encrypted-public", "encrypted-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" not in json.dumps(stored.outer_trace) @@ -607,6 +686,7 @@ def test_source_owner_checks_app_and_tenant(collection: _ApplicationTraceCollect source_kind="message", source_id=run.source_id, session_id="x", + workflow_id=None, ) is None ) @@ -618,19 +698,19 @@ def test_source_owner_checks_app_and_tenant(collection: _ApplicationTraceCollect source_kind="message", source_id=run.source_id, session_id="x", + workflow_id=None, ) 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) attach_conversation(case, AppMode.AGENT) - service = _application_trace_collection_tracing_service(case, _application_trace_collection_no_egress()) - assert case.run.state == "not_collected" + 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 @@ -642,26 +722,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 ) - assert service.capture_terminal(message_record(service, case)) 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) attach_conversation(case, AppMode.AGENT) - service = _application_trace_collection_tracing_service(case, _application_trace_collection_no_egress()) - scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) + ports = ServicePorts(_application_trace_collection_no_egress()) + service = ports.service(case) assert ( service.prepare_invocation( tenant_id=case.run.tenant_id, @@ -673,12 +752,13 @@ def test_capture_agent_disabled_preserves_outer_only_content_policy( node_execution_id=None, attempt=0, call_key="call", + workflow_id=None, ) is None ) - info = message_record(service, case).model_copy(update={"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 @@ -695,7 +775,11 @@ def test_deleting_selected_configuration_keeps_historical_destination( original = collection.run service = AppTracingConfigService( configs=SQLAlchemyAppTracingConfigRepository(session_factory=collection.sessions), - provider=ProviderConfigGateway(_application_trace_collection_no_egress()), + 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) @@ -716,15 +800,14 @@ 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_install_delivery = _application_trace_collection_install_delivery _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 -_application_trace_service_boundaries_tracing_service = _application_trace_collection_tracing_service FROZEN_LANGFUSE: dict[str, object] = { @@ -735,20 +818,39 @@ def test_deleting_selected_configuration_keeps_historical_destination( } -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, - } - - -def forbid_credentials(monkeypatch: pytest.MonkeyPatch, reason: str) -> None: - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail(reason) - - monkeypatch.setattr(provider_configs_module, "batch_decrypt_token", forbidden) +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: @@ -762,13 +864,10 @@ def freeze_config(case: _ApplicationTraceServiceBoundariesCollection, provider: 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 = _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() - ) + 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, @@ -779,20 +878,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 ) @@ -803,29 +906,28 @@ 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: - _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() - ).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 = {**FROZEN_LANGFUSE, "host": "https://new.example", "project_key": "new"} - forbid_credentials(monkeypatch, "status must not decrypt credentials") node_id = _application_trace_service_boundaries_identifier("selected-node") invocation = case.repository.invocation( run_id=case.run.id, @@ -843,33 +945,33 @@ def test_status_uses_frozen_destination_and_selects_latest_node_attempt( case.repository.complete( case.run.id, lease[1], - state="accepted", + state=TraceRunState.ACCEPTED, + error_code=None, result=SendResult( trace_id="provider-trace", trace_url="https://frozen.example/trace", span_urls={invocation.span_id: "https://frozen.example/span"}, ), + retry_delay=None, ) - service = _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() - ) - result = service.status(**source(case), node_execution_id=node_id) - assert result["status"] == "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["queryable"] is None - assert result["capture_content"] is False - capabilities = result["capabilities"] - assert isinstance(capabilities, dict) - assert capabilities["provider"] == "langfuse" - 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 @@ -905,14 +1007,14 @@ def test_new_run_freezes_updated_project_without_changing_previous_run( source_kind="message", source_id=new_id, session_id="new-session", + workflow_id=None, ) assert new_run is not None - service = _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() - ) - assert service.status(**source(case))["project_url"] == "https://langfuse.test/project/original" + service = ServicePorts(_application_trace_service_boundaries_no_egress()).service(case) + 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 @@ -934,10 +1036,12 @@ def test_status_hides_unavailable_or_unsafe_frozen_project_urls( ) -> None: case = _application_trace_service_boundaries_seed(sqlite_session_factory, "invalid-project") freeze_config(case, "langfuse", frozen) - result = _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() - ).status(**source(case)) - assert result["project_url"] is None + 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 @@ -957,23 +1061,25 @@ def test_status_derives_project_link_from_frozen_provider_not_stored_urls( ) -> None: case = _application_trace_service_boundaries_seed(sqlite_session_factory, "platform-console") freeze_config(case, provider, {**frozen, "project_url": "https://stored.example/project"}) - destination = provider_class(provider) + installed = EntryPointTraceProviderCatalog() + destination = installed.resolve(provider) expected = destination(_application_trace_service_boundaries_no_egress()).project_url( destination.config_model.model_validate(frozen) ) - result = _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() - ).status(**source(case)) + 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 - capabilities = result["capabilities"] - assert isinstance(capabilities, dict) - assert capabilities["provider"] == provider + assert result.project_url == expected + assert result.capabilities is not None + assert result.capabilities.provider == provider @staticmethod @pytest.mark.parametrize("disabled", [True, False]) def test_uncollected_status_and_missing_run_retry_do_not_construct_provider( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, disabled: bool + sqlite_session_factory: sessionmaker[Session], disabled: bool ) -> None: case = _application_trace_service_boundaries_seed(sqlite_session_factory, "missing", enabled=not disabled) if not disabled: @@ -981,26 +1087,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) - - def forbidden(_provider_id: str) -> None: - pytest.fail("an uncollected source has no destination to construct") - - monkeypatch.setattr(service_module, "provider_class", forbidden) - service = _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() + 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 ) - assert service.status(**source(case)) == { - "status": "not_collected", - "source_kind": "message", - "source_id": 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( @@ -1009,19 +1109,23 @@ 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 = _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() + case.repository.complete( + case.run.id, + lease[1], + state=TraceRunState.UNCERTAIN, + error_code="provider_acceptance_unknown", + result=None, + retry_delay=None, ) - 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 + 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 @@ -1030,37 +1134,50 @@ def test_retry_acknowledgement_guards_schedule_and_keeps_frozen_snapshot( @pytest.mark.parametrize( ("failure", "max_attempts", "state", "code"), [ - (ExportRetryableError("private-transport-error"), 1, "failed", "provider_retry_exhausted"), - (ExportRetryableError("private-transport-error"), 2, "retrying", "provider_temporarily_unavailable"), - (ExportRejectedError("private-rejection"), 2, "failed", "provider_export_rejected"), - (ExportOutcomeUnknownError("private-acknowledgement"), 2, "uncertain", "provider_acceptance_unknown"), - (ValueError("private-credential-error"), 2, "failed", "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_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: str, + state: TraceRunState, code: str, ) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_MAX_ATTEMPTS=max_attempts) 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) - _application_trace_service_boundaries_install_delivery(monkeypatch) destination = RecordingEgress([failure], grpc_target=None) - service = _application_trace_service_boundaries_tracing_service(case, destination) + 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, current.error_code) == (state, code) - assert (current.next_attempt_at is not None) is (state == "retrying") 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")) @@ -1069,7 +1186,6 @@ def test_delivery_maps_provider_failures_to_stable_codes_without_vendor_details( @pytest.mark.parametrize("unreadable", ["decryption", "validation"]) def test_unreadable_frozen_configuration_fails_before_provider_io( sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, unreadable: str, ) -> None: @@ -1078,28 +1194,22 @@ def test_unreadable_frozen_configuration_fails_before_provider_io( case, _application_trace_service_boundaries_batch_for(case.run, case.invocation()) ) _application_trace_service_boundaries_terminal(case) - _application_trace_service_boundaries_install_delivery(monkeypatch) - if unreadable == "decryption": - - def undecryptable(_tenant_id: str, _tokens: list[str]) -> list[str]: - raise ValueError("private-ciphertext") - - monkeypatch.setattr(provider_configs_module, "batch_decrypt_token", undecryptable) - else: + 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) - _application_trace_service_boundaries_tracing_service(case, destination).deliver(case.run.id) + ServicePorts(destination, cipher=cipher).service(case).deliver(case.run.id) current = case.repository.get(case.run.id) assert current is not None - assert (current.state, current.error_code) == ("failed", "tracing_config_unreadable") + assert (current.state, current.error_code) == (TraceRunState.FAILED, "tracing_config_unreadable") assert destination.requests == [] assert "private-" not in caplog.text - with pytest.raises(TraceCollectionError, match="tracing_retry_not_allowed"): + with pytest.raises(TraceCollectionConflictError, match="tracing_retry_not_allowed"): case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) @staticmethod def test_unavailable_frozen_provider_fails_before_decryption( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + sqlite_session_factory: sessionmaker[Session], ) -> None: case = _application_trace_service_boundaries_seed(sqlite_session_factory, "retired-provider") _application_trace_service_boundaries_ingest( @@ -1107,13 +1217,13 @@ def test_unavailable_frozen_provider_fails_before_decryption( ) _application_trace_service_boundaries_terminal(case) freeze_config(case, "retired-provider", FROZEN_LANGFUSE) - forbid_credentials(monkeypatch, "an unavailable provider must not decrypt credentials") - _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() - ).deliver(case.run.id) + 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) == ("failed", "tracing_provider_unavailable") + assert (current.state, current.error_code) == (TraceRunState.FAILED, "tracing_provider_unavailable") @staticmethod @pytest.mark.parametrize("capture_content", [False, True]) @@ -1166,15 +1276,17 @@ def test_preview_terminal_uses_persisted_node_state_and_frozen_privacy( session_id="preview-session", ) assert run is not None - _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() - ).settle_preview(run) + destination = _application_trace_service_boundaries_accepted("preview-trace") + ServicePorts(destination).service(case).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 + assert current.state is TraceRunState.ACCEPTED trace = Trace.model_validate(current.outer_trace) root = trace.root assert root.span_id == run.root_span_id @@ -1197,7 +1309,7 @@ def test_preview_terminal_uses_persisted_node_state_and_frozen_privacy( ], ) 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()) @@ -1223,27 +1335,28 @@ 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(provider_configs_module, "batch_decrypt_token", forbidden) - monkeypatch.setattr(service_module, "provider_class", forbidden) - _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() - ).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"): - _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() - ).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( @@ -1255,11 +1368,9 @@ def test_assembled_trace_nests_runtime_spans_under_their_invocation_with_trusted case, _application_trace_service_boundaries_batch_for(case.run, invocation) ) _application_trace_service_boundaries_terminal(case) - stored = case.repository.get(case.run.id) - assert stored is not None - trace = _application_trace_service_boundaries_tracing_service( - case, _application_trace_service_boundaries_no_egress() - ).assemble(stored) + 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] == [ @@ -1285,14 +1396,13 @@ def test_assembled_trace_nests_runtime_spans_under_their_invocation_with_trusted _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 -_invalid_trace_delivery_tracing_service = _application_trace_collection_tracing_service 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()) @@ -1316,19 +1426,18 @@ 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(provider_configs_module, "batch_decrypt_token", forbidden) - monkeypatch.setattr(service_module, "provider_class", forbidden) - _invalid_trace_delivery_tracing_service(case, _invalid_trace_delivery_no_egress()).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"): + with pytest.raises(TraceCollectionConflictError, match="tracing_retry_not_allowed"): case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) @@ -1343,7 +1452,6 @@ def forbidden(*_args: object, **_kwargs: object) -> None: _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 -_application_trace_node_alias_tracing_service = _application_trace_collection_tracing_service @dataclass @@ -1353,11 +1461,12 @@ class NodeAlias: workflow: dict[str, object] requests: list[str] logstore: bool + ports: ServicePorts def service(self) -> ApplicationTracingService: - return _application_trace_node_alias_tracing_service(self.case, _application_trace_node_alias_no_egress()) + return self.ports.service(self.case) - def status(self, node_execution_id: str | None = None) -> dict[str, object]: + 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, @@ -1387,12 +1496,14 @@ def accept(self, *, spans: bool = True) -> tuple[str, str]: self.case.repository.complete( self.case.run.id, lease[1], - state="accepted", + state=TraceRunState.ACCEPTED, + error_code=None, result=SendResult( trace_id="provider-trace", trace_url="https://provider.test/trace", span_urls=span_urls if spans else {}, ), + retry_delay=None, ) return "https://provider.test/trace", "https://provider.test/span/1" @@ -1522,6 +1633,7 @@ def execute_sql( source_kind="workflow", source_id=str(workflow["id"]), session_id="caller-session", + workflow_id=None, ) assert run is not None requests.clear() @@ -1531,6 +1643,7 @@ def execute_sql( workflow, requests, request.param == "logstore", + ServicePorts(_application_trace_node_alias_no_egress()), ) finally: event.remove(sqlite_engine, "begin", begun) @@ -1542,13 +1655,13 @@ def execute_sql( 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 + 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"] @@ -1569,18 +1682,20 @@ def test_alias_without_internal_invocation_returns_parent_status( alias.case.repository.complete( alias.case.run.id, lease[1], - state="accepted", + state=TraceRunState.ACCEPTED, + error_code=None, result=SendResult(trace_id="outer-only", trace_url="https://provider.test/outer", span_urls={}), + retry_delay=None, ) 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) + 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", ["collecting", "not_collected"]) - def test_standalone_node_snapshot_wins_over_parent(alias: NodeAlias, state: str) -> None: + @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.case.repository.start( tenant_id=alias.case.run.tenant_id, @@ -1588,6 +1703,7 @@ def test_standalone_node_snapshot_wins_over_parent(alias: NodeAlias, state: str) 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: @@ -1596,15 +1712,13 @@ def test_standalone_node_snapshot_wins_over_parent(alias: NodeAlias, state: str) run.state = state alias.requests.clear() status = alias.status() - assert status["status"] == state - assert status.get("trace_url") is None + 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, monkeypatch: pytest.MonkeyPatch - ) -> None: + def test_explicit_other_node_filter_cannot_select_sibling_span(alias: NodeAlias) -> None: sibling = alias.case.repository.invocation( run_id=alias.case.run.id, agent_id=_application_trace_node_alias_identifier("referenced-agent"), @@ -1616,14 +1730,10 @@ def test_explicit_other_node_filter_cannot_select_sibling_span( alias.case, _application_trace_node_alias_batch_for(alias.case.run, sibling, span_id="f" * 16) ) alias.accept() - with pytest.raises(TraceCollectionError, match="tracing_node_not_found") as error: + with pytest.raises(TraceCollectionNotFoundError, match="tracing_node_not_found"): alias.status(_application_trace_node_alias_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( + 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", @@ -1631,23 +1741,19 @@ def test_explicit_other_node_filter_cannot_select_sibling_span( node_execution_id=_application_trace_node_alias_identifier("sibling-node"), acknowledge_possible_duplicates=True, ) - assert not scheduled + assert not alias.ports.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: + @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() - scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) - def retry(acknowledge: bool) -> dict[str, object]: + def retry(acknowledge: bool) -> TraceRunStatus: alias.requests.clear() return service.retry( tenant_id=alias.case.run.tenant_id, @@ -1655,20 +1761,20 @@ def retry(acknowledge: bool) -> dict[str, object]: source_kind="node", source_id=str(alias.node["id"]), acknowledge_possible_duplicates=acknowledge, + node_execution_id=None, ) - if state == "accepted": - with pytest.raises(TraceCollectionError, match="tracing_already_accepted") as error: + if state is TraceRunState.ACCEPTED: + with pytest.raises(TraceCollectionConflictError, match="tracing_already_accepted"): retry(True) - assert error.value.status == 409 - assert not scheduled + assert not alias.ports.scheduled else: - if state == "uncertain": - with pytest.raises(TraceCollectionError, match="tracing_duplicate_acknowledgement_required"): + if state is TraceRunState.UNCERTAIN: + with pytest.raises(TraceCollectionConflictError, match="tracing_duplicate_acknowledgement_required"): retry(False) - assert not scheduled - assert retry(state == "uncertain")["status"] == "ready" - assert scheduled == [alias.case.run.id] + 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"] @@ -1691,9 +1797,8 @@ def test_alias_rejects_broken_owner_chain(alias: NodeAlias, source: str, field: model.app_id = _application_trace_node_alias_identifier("foreign") else: model.workflow_id = _application_trace_node_alias_identifier("foreign") - with pytest.raises(TraceCollectionError, match="tracing_source_not_found") as error: + with pytest.raises(TraceCollectionNotFoundError, match="tracing_source_not_found"): alias.status() - assert error.value.status == 404 @staticmethod def test_alias_missing_parent_business_source_is_not_a_durable_owner_bypass(alias: NodeAlias) -> None: @@ -1704,7 +1809,7 @@ def test_alias_missing_parent_business_source_is_not_a_durable_owner_bypass(alia 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"): + with pytest.raises(TraceCollectionNotFoundError, match="tracing_source_not_found"): alias.status() @staticmethod @@ -1713,7 +1818,7 @@ def test_owned_alias_without_parent_trace_is_not_collected(alias: NodeAlias) -> 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"]} + assert alias.status() == not_collected("node", str(alias.node["id"])) """Real SQLite state transitions and background renewal for delivery leases. @@ -1728,10 +1833,8 @@ def test_owned_alias_without_parent_trace_is_not_collected(alias: NodeAlias) -> _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_install_delivery = _application_trace_collection_install_delivery _export_lease_seed = _application_trace_collection_seed _export_lease_terminal = _application_trace_collection_terminal -_export_lease_tracing_service = _application_trace_collection_tracing_service class _SlowDestination(RecordingEgress): @@ -1758,14 +1861,52 @@ def request( ) +class _RenewalRecordingRepository(ApplicationTraceRepository): + """Records the thread of every lease renewal; the first one can fail like a dropped connection.""" + + 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 + def renew_export_lease(self, run_id: str, token: str) -> bool: + self.renewal_threads.append(threading.get_ident()) + if self._transient_failure and len(self.renewal_threads) == 1: + raise RuntimeError("private-provider-value") + renewal = super().renew_export_lease(run_id, token) + self.renewed.set() + return renewal + + +class _UnavailableCompletionRepository(ApplicationTraceRepository): + """Loses the database when the export outcome is recorded.""" + + @override + def complete( + self, + run_id: str, + token: str, + *, + state: TraceRunState, + error_code: str | None, + result: SendResult | None, + retry_delay: int | None, + ) -> None: + raise TracingStoreUnavailableError("database unavailable") + + def result(trace_id: str) -> SendResult: return SendResult(trace_id=trace_id, trace_url=None, span_urls={}) def _stored(case: _ExportLeaseCollection) -> ApplicationTraceRun: - row = case.repository.get(case.run.id) - assert row is not None - return row + """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: @@ -1777,7 +1918,7 @@ def _expire(case: _ExportLeaseCollection) -> str: 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 == "uncertain" + assert _stored(case).state == TraceRunState.UNCERTAIN return leased[1] @@ -1804,14 +1945,14 @@ def test_renewal_keeps_long_export_alive_then_missing_heartbeat_expires( clock[0] += timedelta(seconds=20) assert _export_lease_claim(case) is None current = _stored(case) - assert current.state == "exporting" + 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) < clock[0] clock[0] += timedelta(seconds=11) assert _export_lease_claim(case) is None expired = _stored(case) - assert expired.state == "uncertain" + assert expired.state == TraceRunState.UNCERTAIN assert expired.error_code == "export_lease_expired" assert expired.export_token == token assert not case.repository.renew_export_lease(case.run.id, token) @@ -1827,7 +1968,7 @@ def test_old_rows_without_heartbeat_use_started_at(case: _ExportLeaseCollection) row.export_heartbeat_at = None assert _export_lease_claim(case) is None current = _stored(case) - assert current.state == "uncertain" + assert current.state == TraceRunState.UNCERTAIN assert current.export_token == leased[1] @staticmethod @@ -1836,11 +1977,18 @@ def test_wrong_token_cannot_renew_or_complete(case: _ExportLeaseCollection) -> N 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=result("wrong")) + case.repository.complete( + case.run.id, + "other-token", + state=TraceRunState.ACCEPTED, + error_code=None, + result=result("wrong"), + retry_delay=None, + ) current = _stored(case) assert current.export_token == leased[1] assert current.export_heartbeat_at == before.export_heartbeat_at - assert current.state == "exporting" + assert current.state == TraceRunState.EXPORTING assert current.provider_trace_id is None @staticmethod @@ -1851,17 +1999,26 @@ def test_late_success_wins_before_manual_retry_and_is_idempotent(case: _ExportLe trace_url="https://example.com/trace", span_urls={"a" * 16: "https://example.com/span"}, ) - case.repository.complete(case.run.id, token, state="accepted", result=confirmed) + case.repository.complete( + case.run.id, token, state=TraceRunState.ACCEPTED, error_code=None, result=confirmed, retry_delay=None + ) current = _stored(case) - assert current.state == "accepted" + 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"} - case.repository.complete(case.run.id, token, state="accepted", result=result("duplicate-ack")) + case.repository.complete( + case.run.id, + token, + state=TraceRunState.ACCEPTED, + error_code=None, + result=result("duplicate-ack"), + retry_delay=None, + ) assert _stored(case).provider_trace_id == "confirmed" - with pytest.raises(TraceCollectionError, match="tracing_already_accepted"): + with pytest.raises(TraceCollectionConflictError, match="tracing_already_accepted"): case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) @staticmethod @@ -1871,26 +2028,51 @@ def test_manual_retry_revokes_old_heartbeat_and_ack_before_and_after_new_claim( 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=result("old")) - assert _stored(case).state == "ready" + case.repository.complete( + case.run.id, + old_token, + state=TraceRunState.ACCEPTED, + error_code=None, + result=result("old"), + retry_delay=None, + ) + assert _stored(case).state == TraceRunState.READY second = _export_lease_claim(case) assert second is not None assert second[1] != old_token - case.repository.complete(case.run.id, old_token, state="accepted", result=result("old")) + case.repository.complete( + case.run.id, + old_token, + state=TraceRunState.ACCEPTED, + error_code=None, + result=result("old"), + retry_delay=None, + ) current = _stored(case) - assert current.state == "exporting" + assert current.state == TraceRunState.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=result("new")) + case.repository.complete( + case.run.id, + second[1], + state=TraceRunState.ACCEPTED, + error_code=None, + result=result("new"), + retry_delay=None, + ) 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: _ExportLeaseCollection, late_state: str) -> None: + @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) - case.repository.complete(case.run.id, token, state=late_state, error_code="late-failure") + case.repository.complete( + case.run.id, token, state=late_state, error_code="late-failure", result=None, retry_delay=None + ) current = _stored(case) - assert current.state == "uncertain" + assert current.state == TraceRunState.UNCERTAIN assert current.error_code == "export_lease_expired" @staticmethod @@ -1900,16 +2082,18 @@ def test_other_uncertain_reason_cannot_be_promoted_by_late_ack(case: _ExportLeas 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=result("late")) - assert _stored(case).state == "uncertain" + case.repository.complete( + case.run.id, token, state=TraceRunState.ACCEPTED, error_code=None, result=result("late"), retry_delay=None + ) + assert _stored(case).state == TraceRunState.UNCERTAIN @staticmethod @pytest.mark.parametrize("transient_failure", [False, True]) @pytest.mark.parametrize( ("outcome", "reply"), [ - ("accepted", response(200, {"trace_id": "ack"})), - ("uncertain", EgressError("lost acknowledgement", request_sent=True)), + (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( @@ -1917,68 +2101,50 @@ def test_service_heartbeats_during_export_and_stops_on_every_outcome( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, transient_failure: bool, - outcome: str, + outcome: TraceRunState, reply: EgressResponse | Exception, ) -> 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]) - _export_lease_install_delivery(monkeypatch) - renewed = threading.Event() - calls: list[int] = [] - 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") - renewal = original(run_id, token) - renewed.set() - return renewal - - monkeypatch.setattr(case.repository, "renew_export_lease", renew) + 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 = renewed.wait(3) + heartbeat_committed = renewals.renewed.wait(3) second_claim = case.repository.claim(case.run.id, collection_timeout=60, lease_seconds=1) observed.append((heartbeat_committed, second_claim is None, _stored(case).state)) - _export_lease_tracing_service(case, _SlowDestination([reply], while_exporting)).deliver(case.run.id) - assert observed == [(True, True, "exporting")] + 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 calls - assert all(thread_id != threading.get_ident() for thread_id in calls) + 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, monkeypatch: pytest.MonkeyPatch - ) -> None: - _export_lease_install_delivery(monkeypatch) - - 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"): - _export_lease_tracing_service(case, _export_lease_accepted("ack")).deliver(case.run.id) + def test_heartbeat_stops_even_when_completion_database_write_raises(case: _ExportLeaseCollection) -> None: + unavailable = replace(case, repository=_UnavailableCompletionRepository(case.sessions)) + with pytest.raises(TracingStoreUnavailableError, match="database unavailable"): + ServicePorts(_export_lease_accepted("ack")).service(unavailable).deliver(case.run.id) 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 -_application_trace_ownership_tracing_service = _application_trace_collection_tracing_service 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 = _application_trace_ownership_tracing_service(case, RecordingEgress([], grpc_target=None)) + service = ServicePorts(_application_trace_ownership_no_egress()).service(case) record = GenerateNameTraceInfo( message_id=None, message_data=None, @@ -1991,8 +2157,18 @@ def test_naming_without_a_message_stays_with_its_agent_app(sqlite_session_factor tenant_id=case.run.tenant_id, ) - assert service.capture_terminal(record) is True + 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() -> None: - assert ApplicationTracingService._frozen_capabilities("not-a-provider") is None + 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 b96a4e3ee5731b..192abb2428a994 100644 --- a/api/tests/unit_tests/tasks/test_application_trace_task.py +++ b/api/tests/unit_tests/tasks/test_application_trace_task.py @@ -1,48 +1,86 @@ """Exercise Celery entrypoints against durable application trace state.""" -from datetime import timedelta -from uuid import UUID +from dataclasses import replace +from datetime import UTC, datetime, timedelta +from uuid import UUID, uuid4 import httpx import pytest -from sqlalchemy import Engine +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent +from flask import Flask from sqlalchemy.orm import Session, sessionmaker from core.helper import ssrf_proxy -from extensions.ext_database import db +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 models.model import AppMode -from services.tracing import service as service_module -from services.tracing.service import application_tracing_service +from models.enums import ConversationFromSource +from models.model import AppMode, Message +from repositories.trace_record_repository import TraceRecordRepository +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import AgentTraceSettings, TraceRunState +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_no_egress as no_egress -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 ( - _application_trace_collection_tracing_service as tracing_service, +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, ) -from tests.unit_tests.services.tracing.test_service import attach_conversation, install_decryption, message_record -@pytest.fixture -def worker_database(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> None: - """The worker's service factory binds the application engine; bind it to the test database.""" - monkeypatch.setattr(type(db), "engine", property(lambda _db: sqlite_engine)) +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( + traces=collection.repository, + records=TraceRecordRepository(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)), + ) -@pytest.mark.usefixtures("worker_database") 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) - install_decryption(monkeypatch) + bind_worker( + app, + account_application_services, + case, + settings=AGENT_TRACE_SETTINGS, + scheduled=[], + monkeypatch=monkeypatch, + ) proxied: list[tuple[str, str]] = [] def send(method: str, url: str, **_kwargs: object) -> httpx.Response: @@ -55,26 +93,54 @@ def send(method: str, url: str, **_kwargs: object) -> httpx.Response: 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 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}" -@pytest.mark.usefixtures("worker_database") def test_worker_records_carry_no_enterprise_names(sqlite_session_factory: sessionmaker[Session]) -> None: case = seed(sqlite_session_factory, "worker-names") - attach_conversation(case, AppMode.AGENT) - record = message_record(application_tracing_service(), case) + 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") @@ -83,14 +149,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 = tracing_service(due, no_egress()) 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_trace_delivery_task.py b/api/tests/unit_tests/tasks/test_trace_delivery_task.py index befcbeafea0731..9e70a12ba7d04d 100644 --- a/api/tests/unit_tests/tasks/test_trace_delivery_task.py +++ b/api/tests/unit_tests/tasks/test_trace_delivery_task.py @@ -1,15 +1,34 @@ +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 message_trace - -from services.tracing.assembly import SourceNotFoundError, SourceNotReadyError -from services.tracing.delivery import ExportJob, TraceCollectorUnavailableError -from services.tracing.emitter import payload_key +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: @@ -32,19 +51,65 @@ def scan(self, path: str, files: bool = True) -> list[str]: return keys -class _Delivery: - def __init__(self, failure: Exception | None) -> None: - self.failure = failure - self.delivered: list[tuple[TraceEvent, bool]] = [] - self.retried: list[ExportJob] = [] +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") - def deliver(self, event: TraceEvent, *, enterprise_delivered: bool) -> None: - self.delivered.append((event, enterprise_delivered)) - if self.failure is not None: - raise self.failure - def retry(self, job: ExportJob) -> None: - self.retried.append(job) +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: @@ -77,54 +142,67 @@ def storage(monkeypatch: pytest.MonkeyPatch) -> _Storage: return fake -def _use_delivery(monkeypatch: pytest.MonkeyPatch, failure: Exception | None) -> _Delivery: - delivery = _Delivery(failure) - monkeypatch.setattr(trace_delivery_task, "trace_delivery_service", lambda: delivery) +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(monkeypatch: pytest.MonkeyPatch) -> None: - delivery = _use_delivery(monkeypatch, None) - +def test_inline_events_are_delivered(delivery: _Delivery) -> None: deliver_trace_event.run({"event": _event().model_dump_json()}) - ((event, enterprise_delivered),) = delivery.delivered - assert (event.event_id, enterprise_delivered) == ("event-1", False) + assert [event.event_id for event in delivery.records.events] == ["event-1"] -def test_stored_events_are_loaded_and_their_payload_removed(monkeypatch: pytest.MonkeyPatch, storage: _Storage) -> None: - delivery = _use_delivery(monkeypatch, None) +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.delivered) == 1 + assert len(delivery.records.events) == 1 assert key not in storage.files -def test_unreadable_events_are_dropped(monkeypatch: pytest.MonkeyPatch, storage: _Storage) -> None: - delivery = _use_delivery(monkeypatch, None) +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.delivered == [] + assert delivery.records.events == [] assert storage.files == {} @pytest.mark.usefixtures("storage") -def test_missing_sources_are_dropped(monkeypatch: pytest.MonkeyPatch) -> None: - _use_delivery(monkeypatch, SourceNotFoundError()) - +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, storage: _Storage + monkeypatch: pytest.MonkeyPatch, delivery: _Delivery, storage: _Storage ) -> None: - _use_delivery(monkeypatch, SourceNotReadyError()) + delivery.records.result = SourceNotReadyError() retries = _RetryRequests() monkeypatch.setattr(deliver_trace_event, "retry", retries) key = payload_key(datetime.now(UTC), "events", "event-1") @@ -138,8 +216,22 @@ def test_sources_not_ready_are_redelivered_and_keep_their_payload( @pytest.mark.usefixtures("storage") -def test_collector_outages_redeliver_without_repeating_enterprise_telemetry(monkeypatch: pytest.MonkeyPatch) -> None: - _use_delivery(monkeypatch, TraceCollectorUnavailableError()) +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()} @@ -150,25 +242,23 @@ def test_collector_outages_redeliver_without_repeating_enterprise_telemetry(monk assert retries.calls[0]["args"] == ({**envelope, "enterprise_delivered": True},) -def test_export_retries_run_the_stored_job_once(monkeypatch: pytest.MonkeyPatch, storage: _Storage) -> None: - delivery = _use_delivery(monkeypatch, None) +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.retried == [job] + assert delivery.destinations.lookups == [("tenant-1", "app-1")] assert key not in storage.files -def test_unreadable_export_jobs_are_dropped(monkeypatch: pytest.MonkeyPatch, storage: _Storage) -> None: - delivery = _use_delivery(monkeypatch, None) +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.retried == [] + assert delivery.destinations.lookups == [] assert storage.files == {} From 49042f21173a9a76af24823b8870713b12b67d12 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Thu, 29 Oct 2026 00:05:44 +0800 Subject: [PATCH 03/17] refactor(tracing): move Agent trace runs into the domain package Follow the backend architecture guide (#44957). An Agent trace run is now TraceRunAgg in api/domain/tracing, with its invocations, batches and receipts as child entities. Terminal recovery is a domain service of named functions, and domain errors are translated at the application boundary so error codes stay unchanged. The repository loads and saves whole runs under the run's row lock. A domain import contract keeps the package free of application, persistence, framework and legacy kernel code. --- api/.importlinter | 24 +- api/domain/__init__.py | 0 api/domain/tracing/__init__.py | 0 api/domain/tracing/errors.py | 19 + .../tracing/normalization.py | 0 api/domain/tracing/sources.py | 49 + api/domain/tracing/terminal_recovery.py | 65 ++ api/domain/tracing/trace_run_agg.py | 642 +++++++++++++ .../application_services/tracing.py | 5 +- .../application_trace_repository.py | 678 +++++--------- api/services/tracing/collection.py | 256 +---- api/services/tracing/collection_ports.py | 107 +++ api/services/tracing/provider_configs.py | 28 +- api/services/tracing/providers.py | 27 +- api/services/tracing/recovery.py | 80 +- api/services/tracing/service.py | 470 ++++++---- .../console/agent/test_tracing_runs.py | 3 +- .../inner_api/agent/test_tracing.py | 7 +- .../tracing/test_normalization.py | 2 +- .../domain/tracing/test_terminal_recovery.py | 155 +++ .../domain/tracing/test_trace_run_agg.py | 779 +++++++++++++++ .../application_services/test_tracing.py | 6 +- .../test_application_trace_repository.py | 883 ++++++++++++++++-- .../services/tracing/test_provider_configs.py | 8 +- .../services/tracing/test_recovery.py | 134 ++- .../services/tracing/test_runtime_bridge.py | 78 +- .../services/tracing/test_service.py | 405 ++++---- .../tasks/test_application_trace_task.py | 7 +- 28 files changed, 3559 insertions(+), 1358 deletions(-) create mode 100644 api/domain/__init__.py create mode 100644 api/domain/tracing/__init__.py create mode 100644 api/domain/tracing/errors.py rename api/{services => domain}/tracing/normalization.py (100%) create mode 100644 api/domain/tracing/sources.py create mode 100644 api/domain/tracing/terminal_recovery.py create mode 100644 api/domain/tracing/trace_run_agg.py create mode 100644 api/services/tracing/collection_ports.py rename api/tests/unit_tests/{services => domain}/tracing/test_normalization.py (99%) create mode 100644 api/tests/unit_tests/domain/tracing/test_terminal_recovery.py create mode 100644 api/tests/unit_tests/domain/tracing/test_trace_run_agg.py diff --git a/api/.importlinter b/api/.importlinter index 381631cf3f9347..11069a69fbf6b9 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 @@ -1306,10 +1327,10 @@ 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.normalization services.tracing.provider_configs services.tracing.providers services.tracing.records @@ -1352,6 +1373,7 @@ forbidden_modules = 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 diff --git a/api/domain/__init__.py b/api/domain/__init__.py new file mode 100644 index 00000000000000..e69de29bb2d1d6 diff --git a/api/domain/tracing/__init__.py b/api/domain/tracing/__init__.py new file mode 100644 index 00000000000000..e69de29bb2d1d6 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 100% rename from api/services/tracing/normalization.py rename to api/domain/tracing/normalization.py 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..f836154cc726e4 --- /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_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, payload: Mapping[str, Any], *, max_spans: int, max_bytes: int, now: datetime) -> int: + """Append a runtime batch and return its span count; a final batch finalizes its invocation. + + ``payload`` is the batch's validated JSON form. A batch ID replays only + with the same 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(payload, sort_keys=True, separators=(",", ":")).encode() + digest = hashlib.sha256(canonical).hexdigest() + invocation = self.invocations.get(str(payload["invocation_id"])) + if invocation is None: + raise TraceRunNotFoundError("tracing_invocation_not_found") + batch_id = str(payload["batch_id"]) + spans: Sequence[Mapping[str, Any]] = payload["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 = payload["final"] is True + 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: Mapping[str, Any]) -> dict[str, Any]: + stored = dict(span) + 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/extensions/application_services/tracing.py b/api/extensions/application_services/tracing.py index 7944f5c90dc9a9..26e8b3d61ecc9a 100644 --- a/api/extensions/application_services/tracing.py +++ b/api/extensions/application_services/tracing.py @@ -55,8 +55,11 @@ def build_tracing_services(*, database_client: sessionmaker[Session]) -> Tracing 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( - traces=ApplicationTraceRepository(database_client), + 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, files_url=dify_config.FILES_URL, enterprise_details=False), builder=builder, diff --git a/api/repositories/application_trace_repository.py b/api/repositories/application_trace_repository.py index f9b4ab74760245..171fcf8affa225 100644 --- a/api/repositories/application_trace_repository.py +++ b/api/repositories/application_trace_repository.py @@ -1,54 +1,40 @@ -"""Bounded transactions for Agent trace collection and delivery leases. +"""Persistence of Agent trace runs and reads of the sources they are collected for. -Rows become collection records before their session closes. Every database -failure is reported as TracingStoreUnavailableError. +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, override -from uuid import UUID, uuid5 -from dify_agent.protocol.tracing import AgentTraceBatch -from dify_trace_sdk.ids import span_id_for -from dify_trace_sdk.provider import SendResult 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.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 repositories.tracing_store_errors import tracing_store_errors -from services.tracing.collection import ( - CLOSED_STATES, - ApplicationTraceStore, - MessageProgress, - SourceMessage, - SourceNodeExecution, - SourceWorkflowRun, - TraceApp, - TraceAppMode, - TraceBatch, - TraceInvocation, - TraceRun, - TraceRunState, - TraceSourceKind, -) -from services.tracing.errors import ( - TraceCollectionConflictError, - TraceCollectionLimitError, - TraceCollectionNotFoundError, -) -from services.tracing.normalization import safe_attributes +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 = { @@ -89,6 +75,15 @@ def _answer_saved(message: Message) -> bool: 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, @@ -116,6 +111,92 @@ def _run(row: ApplicationTraceRun) -> TraceRun: ) +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 + + def _invocation(row: ApplicationTraceInvocation) -> TraceInvocation: return TraceInvocation( id=row.id, @@ -130,7 +211,7 @@ def _invocation(row: ApplicationTraceInvocation) -> TraceInvocation: ) -class ApplicationTraceRepository(ApplicationTraceStore): +class ApplicationTraceRepository(TraceRunRepository, TraceRunQuery, TraceSourceQuery): def __init__(self, sessions: sessionmaker[Session]) -> None: self._sessions = sessions @@ -288,100 +369,132 @@ def message_progress(self, *, tenant_id: str, app_id: str, message_id: str) -> M ) @override - def start( - self, - *, - tenant_id: str, - app_id: str, - source_kind: TraceSourceKind, - source_id: str, - session_id: str, - workflow_id: str | None, - ) -> TraceRun | None: - # Re-entry uses the already validated, frozen source. It must not depend - # on LogStore availability or indexing latency after collection began. + def app_tracing(self, *, tenant_id: str, app_id: str) -> AppTracingSnapshot | None: with self._read() as session: - if self._owned_app_id(session, tenant_id=tenant_id, app_id=app_id) 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, - ) - ) - if existing is not None: - return _run(existing) - source = self._source_row(tenant_id=tenant_id, app_id=app_id, kind=source_kind, source_id=source_id) - with self._write() as session: - app = session.scalar(select(App).where(App.id == app_id, App.tenant_id == tenant_id)) + app = session.execute(select(App.tracing).where(App.id == app_id, App.tenant_id == tenant_id)).first() 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, - ) - ) - if existing is not None: - return _run(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 - ) - ) - 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 = {} + 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_for(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=TraceRunState.COLLECTING if enabled else TraceRunState.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) ) - return _run(self._insert_or_get(session, 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) -> Generator[LockedTraceRun | None]: + with self._write() as session: + row = session.scalar(select(ApplicationTraceRun).where(ApplicationTraceRun.id == run_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() + 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( @@ -413,327 +526,34 @@ def workflow_run_trace(self, *, tenant_id: str, workflow_run_id: str) -> TraceRu ) return _run(run) if run is not None else None - @override - def record_initialization_failure( - self, *, tenant_id: str, app_id: str, source_kind: TraceSourceKind, source_id: str - ) -> TraceRun | None: - """A missing start snapshot must never fall back to an unmasked exporter.""" - source = self._source_row(tenant_id=tenant_id, app_id=app_id, kind=source_kind, source_id=source_id) - if source is None: - return None - with self._write() as session: - if self._owned_app_id(session, tenant_id=tenant_id, app_id=app_id) 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 _run(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_for(source_id), - session_id=source_id, - capture_content=False, - capture_agent=False, - state=TraceRunState.FAILED, - error_code="tracing_initialization_failed", - terminal_at=naive_utc_now(), - ) - return _run(self._insert_or_get(session, current)) - - @staticmethod - def _insert_or_get(session: Session, record: ApplicationTraceRun) -> ApplicationTraceRun: - """Insert a run; when a concurrent start won the unique source key, return its run instead. - - The first frozen destination wins, so one execution never exports to two providers. - """ - try: - with session.begin_nested(): - session.add(record) - except IntegrityError: - # A locking read sees the committed winner even under MySQL's REPEATABLE READ snapshot. - existing = session.scalar( - select(ApplicationTraceRun) - .where( - ApplicationTraceRun.tenant_id == record.tenant_id, - ApplicationTraceRun.app_id == record.app_id, - ApplicationTraceRun.source_kind == record.source_kind, - ApplicationTraceRun.source_id == record.source_id, - ) - .with_for_update() - ) - if existing is None: - raise - return existing - return record - - @staticmethod - def _locked_run(session: Session, run_id: str) -> ApplicationTraceRun | None: - return session.scalar(select(ApplicationTraceRun).where(ApplicationTraceRun.id == run_id).with_for_update()) - - @override - def invocation( - self, *, run_id: str, agent_id: str, node_execution_id: str | None, attempt: int, call_key: str - ) -> TraceInvocation: - identifier = str(uuid5(UUID(run_id), call_key)) - with self._write() as session: - run = self._locked_run(session, run_id) - if run is None or run.state not in {TraceRunState.COLLECTING, TraceRunState.READY}: - raise TraceCollectionConflictError("tracing_run_not_collecting") - 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 TraceCollectionConflictError("tracing_invocation_conflict") - if invocation is None: - invocation = ApplicationTraceInvocation( - id=identifier, - run_id=run_id, - agent_id=agent_id, - parent_span_id=span_id_for(node_execution_id) if node_execution_id else run.root_span_id, - span_id=span_id_for(f"agent-attempt:{identifier}"), - node_execution_id=node_execution_id, - attempt=attempt, - ) - session.add(invocation) - return _invocation(invocation) - - @override - 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._write() as session: - run = self._locked_run(session, str(batch.run_id)) - if run is None: - raise TraceCollectionNotFoundError("tracing_run_not_found") - if self._owned_app_id(session, tenant_id=run.tenant_id, app_id=run.app_id) is None: - raise TraceCollectionNotFoundError("tracing_owner_not_found") - invocation = session.scalar( - select(ApplicationTraceInvocation).where( - ApplicationTraceInvocation.id == str(batch.invocation_id), - ApplicationTraceInvocation.run_id == run.id, - ) - ) - if invocation is None: - raise TraceCollectionNotFoundError("tracing_invocation_not_found") - 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 TraceCollectionConflictError("tracing_batch_conflict") - return len(previous.spans), run.id - if run.state in CLOSED_STATES or invocation.finalized: - raise TraceCollectionConflictError("tracing_collection_closed") - if any(span.trace_id != run.trace_id for span in batch.spans): - raise TraceCollectionConflictError("tracing_trace_id_mismatch") - if run.span_count + len(batch.spans) > max_spans or run.payload_bytes + len(canonical) > max_bytes: - raise TraceCollectionLimitError("tracing_payload_limit") - 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, - ) - ) - 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 - - @override - def terminal(self, run_id: str, trace: Mapping[str, Any], root_span_id: str) -> None: - with self._write() as session: - run = self._locked_run(session, run_id) - if run is None or run.state in CLOSED_STATES: - 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 = dict(trace) - run.root_span_id = root_span_id - run.terminal_at = run.terminal_at or naive_utc_now() - run.updated_at = naive_utc_now() - @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 - @staticmethod - def _invocation_rows(session: Session, run_id: str) -> list[TraceInvocation]: - return [ - _invocation(row) - for row in session.scalars( - select(ApplicationTraceInvocation) - .where(ApplicationTraceInvocation.run_id == run_id) - .order_by( - ApplicationTraceInvocation.attempt, - ApplicationTraceInvocation.created_at, - ApplicationTraceInvocation.id, - ) - ) - ] - @override def invocations(self, run_id: str) -> list[TraceInvocation]: with self._read() as session: - return self._invocation_rows(session, run_id) + return [_invocation(row) for row in self._invocation_models(session, run_id)] @override - def collected(self, run_id: str) -> tuple[list[TraceInvocation], list[TraceBatch]]: + def batches(self, run_id: str) -> list[TraceBatch]: with self._read() as session: - batches = [ - TraceBatch(invocation_id=row.invocation_id, spans=list(row.spans)) + return [ + TraceBatch( + id=row.id, + invocation_id=row.invocation_id, + payload_hash=row.payload_hash, + spans=list(row.spans), + final=row.final, + ) for row in session.scalars( select(ApplicationTraceBatch) .where(ApplicationTraceBatch.run_id == run_id) .order_by(ApplicationTraceBatch.created_at) ) ] - return self._invocation_rows(session, run_id), batches - - @override - def claim(self, run_id: str, *, collection_timeout: int, lease_seconds: int) -> tuple[TraceRun, str] | None: - now = naive_utc_now() - with self._write() as session: - run = self._locked_run(session, run_id) - if run is None or run.state in { - TraceRunState.NOT_COLLECTED, - TraceRunState.ACCEPTED, - TraceRunState.FAILED, - TraceRunState.UNCERTAIN, - }: - return None - if run.state == TraceRunState.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 = TraceRunState.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 = TraceRunState.FAILED, "agent_trace_collection_incomplete" - run.updated_at = now - return None - if run.next_attempt_at and run.next_attempt_at > now: - return None - if self._owned_app_id(session, tenant_id=run.tenant_id, app_id=run.app_id) is None: - run.state, run.error_code = TraceRunState.FAILED, "tracing_owner_not_found" - return None - token = str(uuidv7()) - run.state = TraceRunState.EXPORTING - run.export_token = token - run.export_started_at = now - run.export_heartbeat_at = now - run.updated_at = now - run.attempts += 1 - return _run(run), token - - @override - def renew_export_lease(self, run_id: str, token: str) -> bool: - """Renew only the current export attempt, in a short independent transaction.""" - with self._write() as session: - run = self._locked_run(session, run_id) - if run is None or run.state != TraceRunState.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 - - @override - def complete( - self, - run_id: str, - token: str, - *, - state: TraceRunState, - error_code: str | None, - result: SendResult | None, - retry_delay: int | None, - ) -> None: - with self._write() as session: - run = self._locked_run(session, run_id) - if run is None or run.export_token != token: - return - late_acceptance = ( - state is TraceRunState.ACCEPTED - and result is not None - and run.state == TraceRunState.UNCERTAIN - and run.error_code == "export_lease_expired" - ) - if run.state != TraceRunState.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, - dict(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 - - @override - def recovery_checked(self, run_id: str, *, error_code: str | None) -> None: - """Rotate checked incomplete work without changing a concurrent export.""" - with self._write() as session: - run = self._locked_run(session, run_id) - if ( - run is None - or run.outer_trace is not None - or run.state not in {TraceRunState.COLLECTING, TraceRunState.READY} - ): - return - run.updated_at = naive_utc_now() - if error_code is not None: - run.state, run.error_code = TraceRunState.FAILED, error_code @override def pending_ids(self, limit: int) -> list[str]: @@ -760,32 +580,6 @@ def pending_ids(self, limit: int) -> list[str]: ) ) - @override - def retry(self, *, run_id: str, acknowledge_possible_duplicates: bool) -> None: - with self._write() as session: - run = self._locked_run(session, run_id) - if run is None or run.state == TraceRunState.NOT_COLLECTED: - raise TraceCollectionNotFoundError("tracing_run_not_found") - if run.state == TraceRunState.ACCEPTED: - raise TraceCollectionConflictError("tracing_already_accepted") - if run.state not in {TraceRunState.FAILED, TraceRunState.UNCERTAIN} or run.outer_trace is None: - raise TraceCollectionConflictError("tracing_retry_not_allowed") - 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 TraceCollectionConflictError("tracing_retry_not_allowed") - if run.state == TraceRunState.UNCERTAIN and not acknowledge_possible_duplicates: - raise TraceCollectionConflictError("tracing_duplicate_acknowledgement_required") - run.state, run.error_code, run.export_token = TraceRunState.READY, None, None - run.next_attempt_at = None - run.updated_at = naive_utc_now() - @override def purge(self, *, retention_days: int, limit: int) -> int: cutoff = naive_utc_now() - timedelta(days=retention_days) diff --git a/api/services/tracing/collection.py b/api/services/tracing/collection.py index ebb8a71ea08561..5b31a216d59b51 100644 --- a/api/services/tracing/collection.py +++ b/api/services/tracing/collection.py @@ -1,29 +1,22 @@ -"""Contracts of Agent trace collection: run records, the ledger port, and its settings. +"""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. - -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. """ import json -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Mapping from dataclasses import dataclass from datetime import datetime from enum import StrEnum -from typing import Any, Literal, Protocol - -from dify_agent.protocol.tracing import AgentTraceBatch -from dify_trace_sdk.provider import SendResult +from typing import Any from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data -from services.tracing.provider_configs import TraceProviderSupport - -type TraceSourceKind = Literal["message", "workflow", "node"] +from domain.tracing.sources import SourceNodeExecution, SourceWorkflowRun +from domain.tracing.trace_run_agg import TraceRunState, TraceSourceKind +from services.tracing.providers import TraceProviderSupport def workflow_has_agent_v2(graph: object) -> bool: @@ -50,26 +43,6 @@ class TraceAppMode(StrEnum): OTHER = "other" -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.""" - - -# A closed run accepts no runtime batches and no other business terminal. -CLOSED_STATES = frozenset( - {TraceRunState.ACCEPTED, TraceRunState.EXPORTING, TraceRunState.UNCERTAIN, TraceRunState.FAILED} -) - - @dataclass(frozen=True, slots=True) class TraceRun: id: str @@ -104,27 +77,6 @@ class TraceRun: """``None`` until the destination accepts the trace.""" -@dataclass(frozen=True, slots=True) -class TraceInvocation: - 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 TraceBatch: - invocation_id: str - spans: Sequence[Mapping[str, Any]] - - @dataclass(frozen=True, slots=True) class TraceApp: id: str @@ -144,52 +96,20 @@ class SourceMessage: """Whether the consumer asked to stop the Agent; it does not say whether the stop won.""" -@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 - - type TraceSource = SourceMessage | SourceWorkflowRun | SourceNodeExecution @dataclass(frozen=True, slots=True) -class MessageProgress: - """What recovery may infer about a message from its committed row.""" +class AppTracingSnapshot: + """What a run freezes from its app: the stored tracing policy and the provider configuration it names.""" - 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 + 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) @@ -230,149 +150,3 @@ class AgentTraceSettings: export_lease_seconds: int max_export_attempts: int retention_days: int - - -type ScheduleTraceExport = Callable[[str], object] - - -class ApplicationTraceStore(Protocol): - """The Agent trace ledger. Reads and writes are scoped by the tenant and app that own the run. - - Every operation raises TracingStoreUnavailableError when the database - cannot be reached, and the TraceCollectionError subtypes it names. - """ - - def app(self, *, tenant_id: str, app_id: str) -> TraceApp | None: ... - - 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: ... - - def start( - self, - *, - tenant_id: str, - app_id: str, - source_kind: TraceSourceKind, - source_id: str, - session_id: str, - workflow_id: str | None, - ) -> TraceRun | None: - """Freeze the app's destination and capture policy for a source, once. - - A run collects only when the app's enabled provider has a stored - configuration; otherwise it is frozen as not collected. The unique source - key keeps the first frozen run when executions race. ``workflow_id`` lets - a trusted workflow builder reserve a run before its source row is saved. - ``None`` when the source or its owner does not exist. - """ - ... - - 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 record_initialization_failure( - 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.""" - ... - - def invocation( - self, *, run_id: str, agent_id: str, node_execution_id: str | None, attempt: int, call_key: str - ) -> TraceInvocation: - """Create the invocation of a call key once, under the run's row lock. - - Raises: - TraceCollectionConflictError: The run stopped collecting, or the call key names another invocation. - """ - ... - - def ingest(self, batch: AgentTraceBatch, *, max_spans: int, max_bytes: int) -> tuple[int, str]: - """Append a runtime batch under the run's row lock and return its span count and run ID. - - A batch ID replays only with the same payload hash; a final batch finalizes its invocation. - - Raises: - TraceCollectionNotFoundError: The run, its owner, or the invocation does not exist. - TraceCollectionConflictError: The batch conflicts with the ledger or the run stopped collecting. - TraceCollectionLimitError: The batch exceeds the run's limits. - """ - ... - - def terminal(self, run_id: str, trace: Mapping[str, Any], root_span_id: str) -> None: - """Persist the business terminal once under the run's row lock; a different later terminal is ignored.""" - ... - - def get(self, run_id: str) -> TraceRun | None: ... - - def invocations(self, run_id: str) -> list[TraceInvocation]: - """The run's invocations by attempt, then creation.""" - ... - - def collected(self, run_id: str) -> tuple[list[TraceInvocation], list[TraceBatch]]: - """The run's invocations by attempt, then creation, and its batches by arrival.""" - ... - - def claim(self, run_id: str, *, collection_timeout: int, lease_seconds: int) -> tuple[TraceRun, str] | None: - """Lease a complete run for one export attempt and return it with the attempt's token. - - 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. The run's row lock - serializes claims; the returned token fences the attempt. - """ - ... - - def renew_export_lease(self, run_id: str, token: str) -> bool: ... - - def complete( - self, - run_id: str, - token: str, - *, - state: TraceRunState, - error_code: str | None, - result: SendResult | None, - retry_delay: int | None, - ) -> None: - """Record an export outcome under the run's row lock. - - An attempt whose token was replaced records nothing. An acceptance that - arrives after its lease expired still records the destination's result. - """ - ... - - def recovery_checked(self, run_id: str, *, error_code: str | None) -> None: ... - - def pending_ids(self, limit: int) -> list[str]: ... - - def retry(self, *, run_id: str, acknowledge_possible_duplicates: bool) -> None: - """Release a failed or uncertain run for another export. - - Raises: - TraceCollectionNotFoundError: Nothing was collected. - TraceCollectionConflictError: The run is accepted, not retryable, or needs the duplicate acknowledgement. - """ - ... - - def purge(self, *, retention_days: int, limit: int) -> int: ... diff --git a/api/services/tracing/collection_ports.py b/api/services/tracing/collection_ports.py new file mode 100644 index 00000000000000..d25fedd92d6377 --- /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) -> 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. 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/provider_configs.py b/api/services/tracing/provider_configs.py index 347c0031201dd0..6623dd4a442130 100644 --- a/api/services/tracing/provider_configs.py +++ b/api/services/tracing/provider_configs.py @@ -9,7 +9,7 @@ import logging from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Literal, override +from typing import Any, override from dify_trace_sdk.config import ProviderConfig from dify_trace_sdk.egress import Egress @@ -35,7 +35,9 @@ SUPPORTED_PROVIDERS, ProviderConfigCodec, TraceProviderCatalog, + TraceProviderSupport, configured_secrets, + support_of, ) logger = logging.getLogger(__name__) @@ -175,21 +177,6 @@ def _verified( return self._codec.encode(workspace_id, verified) -@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.""" - - @dataclass(frozen=True, slots=True) class TraceProviderCapabilities: support: TraceProviderSupport @@ -197,15 +184,6 @@ class TraceProviderCapabilities: """Whether this deployment collects Agent runtime spans; it does not probe the Agent runtime.""" -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, - ) - - class TraceProviderCapabilitiesQuery: def __init__(self, *, providers: TraceProviderCatalog, api_collection_enabled: bool) -> None: self._providers = providers diff --git a/api/services/tracing/providers.py b/api/services/tracing/providers.py index b45ef871aa0fbe..d7d58385afa58a 100644 --- a/api/services/tracing/providers.py +++ b/api/services/tracing/providers.py @@ -8,7 +8,8 @@ import logging from collections.abc import Mapping, Sequence -from typing import Any, Protocol +from dataclasses import dataclass +from typing import Any, Literal, Protocol from dify_trace_sdk.config import ProviderConfig from dify_trace_sdk.provider import TraceProvider @@ -102,3 +103,27 @@ def mask(self, value: str) -> str: 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/recovery.py b/api/services/tracing/recovery.py index de02588f35c247..18e69ff0c16234 100644 --- a/api/services/tracing/recovery.py +++ b/api/services/tracing/recovery.py @@ -6,13 +6,16 @@ """ import logging -from datetime import UTC, datetime, timedelta +from datetime import UTC, datetime from uuid import uuid4 from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent, TracePayload, WorkflowRunCompleted -from graphon.enums import WorkflowExecutionStatus -from services.tracing.collection import ApplicationTraceStore, TraceRun, TraceRunState +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 services.tracing.collection_ports import TraceRunRepository, TraceSourceQuery from services.tracing.delivery import TraceCollector, TraceRecordSource from services.tracing.errors import ( SourceNotFoundError, @@ -24,16 +27,9 @@ 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, -} -def _event(run: TraceRun, payload: TracePayload, user_id: str | None) -> TraceEvent: +def _event(run: TraceRunAgg, payload: TracePayload, user_id: str | None) -> TraceEvent: return TraceEvent( event_id=uuid4().hex, tenant_id=run.tenant_id, @@ -51,49 +47,22 @@ def _event(run: TraceRun, payload: TracePayload, user_id: str | None) -> TraceEv def _message_event( - traces: ApplicationTraceStore, run: TraceRun, now: datetime, collection_timeout: int + sources: TraceSourceQuery, run: TraceRunAgg, now: datetime, collection_timeout: int ) -> TraceEvent | None: - progress = traces.message_progress(tenant_id=run.tenant_id, app_id=run.app_id, message_id=run.source_id) + 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") - pending_agent_stop = progress.agent_app and progress.agent_stop_requested - if progress.paused and not pending_agent_stop: - return None - invocations = traces.invocations(run.id) - 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 TraceCollectionError("agent_trace_collection_incomplete") - return None - 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 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 - if not (progress.explicit_error or progress.answer_saved): - if finalized and max(finalized) + timedelta(seconds=collection_timeout) < now: - raise TraceCollectionError("business_terminal_not_persisted") + 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_event(traces: ApplicationTraceStore, run: TraceRun, now: datetime) -> TraceEvent | None: - workflow = traces.workflow_run(tenant_id=run.tenant_id, app_id=run.app_id, workflow_run_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 workflow is None or workflow.status not in _WORKFLOW_TERMINAL or workflow.finished_at is None: - return None - if workflow.finished_at + timedelta(seconds=_SETTLE_SECONDS) > now: + if workflow is None or not workflow_terminal_due(workflow, now=now): return None - conversation_id = traces.workflow_run_conversation_id( + conversation_id = sources.workflow_run_conversation_id( tenant_id=run.tenant_id, app_id=run.app_id, workflow_run_id=workflow.id ) return _event( @@ -104,9 +73,10 @@ def _workflow_event(traces: ApplicationTraceStore, run: TraceRun, now: datetime) def recover_terminal( - run: TraceRun, + run: TraceRunAgg, *, - traces: ApplicationTraceStore, + runs: TraceRunRepository, + sources: TraceSourceQuery, records: TraceRecordSource, collector: TraceCollector, collection_timeout: int, @@ -116,18 +86,14 @@ def recover_terminal( ``now`` is naive UTC, like the ledger's timestamps. """ - if ( - run.outer_trace is not None - or run.state not in {TraceRunState.COLLECTING, TraceRunState.READY} - or run.source_kind not in {"message", "workflow"} - ): + if not run.awaits_terminal or run.source_kind not in {"message", "workflow"}: return False error_code: str | None = None try: event = ( - _message_event(traces, run, now, collection_timeout) + _message_event(sources, run, now, collection_timeout) if run.source_kind == "message" - else _workflow_event(traces, run, now) + else _workflow_event(sources, run, now) ) if event is None: return False @@ -142,7 +108,7 @@ def recover_terminal( 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 TracingStoreUnavailableError: @@ -154,4 +120,6 @@ def recover_terminal( logger.warning("Application trace recovery failed: run_id=%s error_type=%s", run.id, type(error).__name__) return False finally: - traces.recovery_checked(run.id, error_code=error_code) + with runs.change(run.id) as locked: + if locked is not None: + locked.run.recovery_checked(error_code, now=naive_utc_now()) diff --git a/api/services/tracing/service.py b/api/services/tracing/service.py index ced9a3d47e6689..39c1a8fcfd827a 100644 --- a/api/services/tracing/service.py +++ b/api/services/tracing/service.py @@ -9,9 +9,9 @@ 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 datetime import UTC, datetime from typing import Any, Literal, override from uuid import uuid4 @@ -20,46 +20,37 @@ 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.ids import span_id_for -from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace 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 ( - CLOSED_STATES, AgentTraceSettings, - ApplicationTraceStore, - ScheduleTraceExport, SourceMessage, - SourceNodeExecution, TraceApp, TraceAppMode, TraceRun, - TraceRunState, TraceRunStatus, TraceSource, - TraceSourceKind, 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.normalization import ( - legacy_outer_trace, - normalize_agent_span, - order_spans, - remap_outer_trace, - sanitize, -) -from services.tracing.provider_configs import TraceProviderSupport, support_of -from services.tracing.providers import ProviderConfigCodec, TraceProviderCatalog +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 @@ -79,11 +70,55 @@ def _message_inputs_are_empty(inputs: object) -> bool: return inputs is None or isinstance(inputs, (str, dict, list)) and not inputs +@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, *, - traces: ApplicationTraceStore, + runs: TraceRunRepository, + run_query: TraceRunQuery, + sources: TraceSourceQuery, records: TraceRecordSource, builder: CanonicalTraceBuilder, providers: TraceProviderCatalog, @@ -92,7 +127,9 @@ def __init__( schedule_export: ScheduleTraceExport, settings: AgentTraceSettings, ) -> None: - self._traces = traces + self._runs = runs + self._run_query = run_query + self._sources = sources self._records = records self._builder = builder self._providers = providers @@ -111,10 +148,13 @@ def start( session_id: str, workflow_id: str | None, ) -> TraceRun | None: - """Freeze collection for a source; ``None`` when the source is not collected.""" + """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._traces.start( + run = self._freeze( tenant_id=tenant_id, app_id=app_id, source_kind=source_kind, @@ -124,6 +164,83 @@ def start( ) 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, *, @@ -148,9 +265,12 @@ def prepare_invocation( ) if run is None or not run.capture_agent: return None - invocation = self._traces.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) 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, @@ -161,9 +281,25 @@ def prepare_invocation( ) def ingest(self, batch: AgentTraceBatch) -> int: - accepted, run_id = self._traces.ingest( - batch, max_spans=self._settings.max_spans, max_bytes=self._settings.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) + with _collection_errors(), self._runs.change(run_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.model_dump(mode="json"), + 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) return accepted @@ -183,10 +319,12 @@ 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._traces.message(tenant_id=tenant_id, app_id=app_id, message_id=message_id) + 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._traces.source_run(tenant_id=tenant_id, app_id=app_id, source_kind="message", source_id=message_id) + 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 is TraceRunState.NOT_COLLECTED: return info = self._records.assemble( @@ -229,9 +367,9 @@ def capture_agent_app_terminal( def reconcile(self) -> None: """Schedule unfinished runs again, behind a broker outage, and purge runs past their retention.""" - for run_id in self._traces.pending_ids(_RECONCILE_BATCH): + for run_id in self._run_query.pending_ids(_RECONCILE_BATCH): self._schedule(run_id) - self._traces.purge(retention_days=self._settings.retention_days, limit=_RECONCILE_BATCH) + self._runs.purge(retention_days=self._settings.retention_days, limit=_RECONCILE_BATCH) def _schedule(self, run_id: str) -> None: try: @@ -239,14 +377,22 @@ def _schedule(self, run_id: str) -> None: except Exception: logger.warning("Application trace is persisted; scheduling deferred: run_id=%s", run_id) + def _change(self, run_id: str, 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) as locked: + if locked is not None: + change(locked.run) + def _run_for_source( self, *, tenant_id: str, app_id: str, source_kind: TraceSourceKind, source_id: str ) -> TraceRun | None: - run = self._traces.source_run(tenant_id=tenant_id, app_id=app_id, source_kind=source_kind, source_id=source_id) + 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": - message = self._traces.message(tenant_id=tenant_id, app_id=app_id, message_id=source_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._traces.source_run( + 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 @@ -264,11 +410,11 @@ def _managed_owner( if app.mode is TraceAppMode.AGENT: return True if workflow_run_id: - if self._traces.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._traces.workflow_run( + workflow = self._sources.workflow_run( tenant_id=app.tenant_id, app_id=app.id, workflow_run_id=workflow_run_id ) if workflow is not None: @@ -278,7 +424,7 @@ def _managed_owner( # the owned app workflow; an unknown snapshot fails closed. The same # applies to a Workflow whose run is not yet indexed by LogStore. graph = ( - self._traces.app_workflow_graph(tenant_id=app.tenant_id, app_id=app.id, workflow_id=app.workflow_id) + 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 ) @@ -291,10 +437,10 @@ def link_managed_parent(self, *, tenant_id: str, info: WorkflowTraceInfo) -> Wor parent_run_id, parent_node_id = info.resolved_parent_context if not parent_run_id: return info - parent = self._traces.workflow_run_trace(tenant_id=tenant_id, workflow_run_id=parent_run_id) + parent = self._run_query.workflow_run_trace(tenant_id=tenant_id, workflow_run_id=parent_run_id) if parent is None: return info - if self._traces.workflow_run(tenant_id=tenant_id, app_id=parent.app_id, workflow_run_id=parent_run_id) is None: + 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) @@ -321,16 +467,16 @@ def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: 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._traces.app(tenant_id=tenant_id, app_id=owner_app_id) + 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._traces.app(tenant_id=tenant_id, app_id=app_id) + app = self._sources.app(tenant_id=tenant_id, app_id=app_id) if app is None: return False message = ( - self._traces.message(tenant_id=tenant_id, app_id=app_id, message_id=source_id) + self._sources.message(tenant_id=tenant_id, app_id=app_id, message_id=source_id) if source_kind == "message" else None ) @@ -342,11 +488,8 @@ def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: if run is None: if not managed: return False - run = self._traces.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( @@ -354,7 +497,7 @@ def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: ) if run is None: return True - if run.state is TraceRunState.NOT_COLLECTED or run.state in CLOSED_STATES: + if not run.state.accepts_terminal: return True if ( app.mode is TraceAppMode.AGENT @@ -369,120 +512,33 @@ def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: trace = self._builder.build(info) if trace is None: raise TraceCollectionError("tracing_terminal_not_supported") - synthetic = frozenset({span_id_for(f"{source_id}:llm")}) if source_kind == "message" else frozenset() - trace = remap_outer_trace( - trace, trace_id=run.trace_id, capture_content=run.capture_content, synthetic_span_ids=synthetic - ) - trace = trace.model_copy(update={"session_id": run.session_id}) - self._traces.terminal(run.id, trace.model_dump(mode="json"), trace.root.span_id) + self._change(run.id, lambda current: current.capture_terminal(trace, now=naive_utc_now())) self._schedule(run.id) return True - def _settle_preview(self, run: TraceRun) -> None: - if run.source_kind != "node" or run.outer_trace is not None: - return - node = self._traces.node_execution(tenant_id=run.tenant_id, app_id=run.app_id, node_execution_id=run.source_id) - if node is None or 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 = Span( - span_id=run.root_span_id, - parent_span_id=None, - name=node.title or "Agent preview", - kind=SpanKind.AGENT, - start_time=started_at, - end_time=ended_at, - status=SpanStatus.ERROR if error else SpanStatus.OK, - error=error if run.capture_content else "Node execution failed" if error else None, - inputs=sanitize(node.inputs) if run.capture_content else None, - outputs=sanitize(node.outputs) if run.capture_content else None, - attributes={"node_execution_id": node.id, "node_id": node.node_id, "source": "single_step"}, - ) - trace = Trace( - trace_id=run.trace_id, - external_trace_id=None, - session_id=run.session_id, - user_id=None, - parent=None, - spans=(span,), - ) - self._traces.terminal(run.id, trace.model_dump(mode="json"), run.root_span_id) - - def _assemble(self, run: TraceRun) -> Trace: - if run.outer_trace is None: - raise TraceCollectionError("tracing_terminal_missing") - # Runs captured before the upgrade stored the previous trace shape, named by its root span. - outer = ( - legacy_outer_trace(run.outer_trace) - if "root_span_id" in run.outer_trace - else Trace.model_validate(run.outer_trace) - ) - invocations, batches = self._traces.collected(run.id) - invocation_by_id = {inv.id: inv for inv in invocations} - spans = list(outer.spans) - internal_by_invocation: dict[str, dict[str, Span]] = {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.span_id) - if current is not None and current != span: - raise TraceCollectionError("tracing_span_conflict") - internal_by_invocation[batch.invocation_id][span.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.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 + 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.id, lambda current: current.settle_preview(node, now=naive_utc_now())) + + def _claim(self, run_id: str) -> tuple[TraceRunAgg, str] | None: + token = str(uuidv7()) + with self._runs.change(run_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(), ): - raise TraceCollectionError("tracing_invocation_parent_invalid") - root_children = [child for child in children if child.parent_span_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 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": 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={"attributes": {**span.attributes, **trusted}}) for span in spans] - return Trace( - trace_id=run.trace_id, - external_trace_id=None, - session_id=run.session_id, - user_id=outer.user_id, - parent=None, - spans=order_spans(outer.root.span_id, spans), - ) + return None + return locked.run, token + + def _renew_export_lease(self, run_id: str, token: str) -> bool: + with self._runs.change(run_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]: @@ -498,7 +554,7 @@ def _export_lease_heartbeat(self, run_id: str, token: str) -> Generator[None]: def heartbeat() -> None: while not stopped.wait(interval): try: - if not self._traces.renew_export_lease(run_id, token): + if not self._renew_export_lease(run_id, token): return except Exception as error: logger.warning( @@ -526,59 +582,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._traces.get(run_id) - if record is None: + run = self._runs.load(run_id) + if run is None: return recover_terminal( - record, - traces=self._traces, + run, + runs=self._runs, + sources=self._sources, records=self._records, collector=self, collection_timeout=self._settings.collection_timeout_seconds, now=naive_utc_now(), ) - record = self._traces.get(run_id) - if record is None: + run = self._runs.load(run_id) + if run is None: return - self._settle_preview(record) - claim = self._traces.claim( - run_id, - collection_timeout=self._settings.collection_timeout_seconds, - lease_seconds=self._settings.export_lease_seconds, - ) + self._settle_preview(run) + claim = self._claim(run_id) if claim is None: return run, token = claim with self._export_lease_heartbeat(run_id, token): try: try: - trace = self._assemble(run) + 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 - provider = self._frozen_provider(run) + provider = self._frozen_provider(run.provider) result = provider.send(self._codec.decode(run.tenant_id, type(provider), run.provider_config), trace) except ExportOutcomeUnknownError: - self._complete_failed(run_id, token, TraceRunState.UNCERTAIN, "provider_acceptance_unknown") + self._change(run_id, lambda current: current.export_uncertain(token, now=naive_utc_now())) except ExportRetryableError: - if run.attempts >= self._settings.max_export_attempts: - self._complete_failed(run_id, token, TraceRunState.FAILED, "provider_retry_exhausted") - else: - self._traces.complete( - run_id, - token, - state=TraceRunState.RETRYING, - error_code="provider_temporarily_unavailable", - result=None, - retry_delay=min(300, 5 * 2 ** min(run.attempts - 1, 6)), - ) - except TraceCollectionError as error: - self._complete_failed(run_id, token, TraceRunState.FAILED, error.code) + self._change( + run_id, + 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_id, token, error.code) except ProviderConfigUnreadableError: - self._complete_failed(run_id, token, TraceRunState.FAILED, "tracing_config_unreadable") + self._export_failed(run_id, token, "tracing_config_unreadable") except ExportRejectedError: - self._complete_failed(run_id, token, TraceRunState.FAILED, "provider_export_rejected") + self._export_failed(run_id, token, "provider_export_rejected") except Exception as error: logger.warning( "Application trace export failed: run_id=%s provider=%s error_type=%s", @@ -586,18 +634,25 @@ def deliver(self, run_id: str) -> None: run.provider, type(error).__name__, ) - self._complete_failed(run_id, token, TraceRunState.FAILED, "provider_export_rejected") + self._export_failed(run_id, token, "provider_export_rejected") else: - self._traces.complete( - run_id, token, state=TraceRunState.ACCEPTED, error_code=None, result=result, retry_delay=None + self._change( + run_id, + 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 _complete_failed(self, run_id: str, token: str, state: TraceRunState, error_code: str) -> None: - self._traces.complete(run_id, token, state=state, error_code=error_code, result=None, retry_delay=None) + def _export_failed(self, run_id: str, token: str, error_code: str) -> None: + self._change(run_id, lambda current: current.export_failed(token, error_code, now=naive_utc_now())) - def _frozen_provider(self, run: TraceRun) -> TraceProvider[Any]: + def _frozen_provider(self, provider_id: str) -> TraceProvider[Any]: try: - return self._providers.resolve(run.provider)(self._egress) + return self._providers.resolve(provider_id)(self._egress) except (UnknownTraceProviderError, TraceProviderUnavailableError): raise TraceCollectionError("tracing_provider_unavailable") from None @@ -636,11 +691,11 @@ def status( def _source(self, *, tenant_id: str, app_id: str, kind: TraceSourceKind, source_id: str) -> TraceSource | None: match kind: case "message": - return self._traces.message(tenant_id=tenant_id, app_id=app_id, message_id=source_id) + return self._sources.message(tenant_id=tenant_id, app_id=app_id, message_id=source_id) case "workflow": - return self._traces.workflow_run(tenant_id=tenant_id, app_id=app_id, workflow_run_id=source_id) + return self._sources.workflow_run(tenant_id=tenant_id, app_id=app_id, workflow_run_id=source_id) case "node": - return self._traces.node_execution(tenant_id=tenant_id, app_id=app_id, node_execution_id=source_id) + return self._sources.node_execution(tenant_id=tenant_id, app_id=app_id, node_execution_id=source_id) def _resolve_status_source( self, @@ -661,12 +716,12 @@ def _resolve_status_source( # 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, SourceNodeExecution) and source.workflow_run_id: - workflow = self._traces.workflow_run( + workflow = self._sources.workflow_run( tenant_id=tenant_id, app_id=app_id, workflow_run_id=source.workflow_run_id ) if workflow is None or not source.workflow_id or workflow.workflow_id != source.workflow_id: raise TraceCollectionNotFoundError("tracing_source_not_found") - run = self._traces.source_run( + 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 @@ -694,7 +749,7 @@ def _status_for_run( accepted_at=None, capabilities=None, ) - invocations = self._traces.invocations(run.id) + 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 [] ) @@ -739,6 +794,11 @@ def retry( if run is None: raise TraceCollectionNotFoundError("tracing_run_not_found") self._status_for_run(run, source_kind, source_id, node_execution_id) - self._traces.retry(run_id=run.id, acknowledge_possible_duplicates=acknowledge_possible_duplicates) + with _collection_errors(), self._runs.change(run.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._traces.get(run.id), source_kind, source_id, node_execution_id) + return self._status_for_run(self._run_query.get(run.id), source_kind, source_id, node_execution_id) 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 a5d58799475dba..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 @@ -16,13 +16,14 @@ 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 services.app.agent_app_contracts import AgentAppNotFoundError -from services.tracing.collection import TraceRunState, TraceRunStatus, TraceSourceKind +from services.tracing.collection import TraceRunStatus from services.tracing.errors import TraceCollectionConflictError, TraceCollectionError, TraceCollectionNotFoundError from tests.helpers.unit.config_override import apply_config_overrides 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 65100cd5cfbb2c..2def94dddb1b3f 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 @@ -15,7 +15,8 @@ from models.application_trace import ApplicationTraceBatch from repositories.trace_record_repository import TraceRecordRepository from services.tracing.canonical import CanonicalTraceBuilder -from services.tracing.collection import AgentTraceSettings, ScheduleTraceExport +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 @@ -49,7 +50,9 @@ def bind_collector( ) -> None: """Make the app's composition root hand the ingestion handler a collector over the seeded ledger.""" collector = ApplicationTracingService( - traces=collection.repository, + runs=collection.repository, + run_query=collection.repository, + sources=collection.repository, records=TraceRecordRepository(collection.sessions, files_url="https://files.test", enterprise_details=False), builder=CanonicalTraceBuilder(), providers=InstalledTraceProviders(), diff --git a/api/tests/unit_tests/services/tracing/test_normalization.py b/api/tests/unit_tests/domain/tracing/test_normalization.py similarity index 99% rename from api/tests/unit_tests/services/tracing/test_normalization.py rename to api/tests/unit_tests/domain/tracing/test_normalization.py index f3722350b6ceda..fa74baab3ad8a4 100644 --- a/api/tests/unit_tests/services/tracing/test_normalization.py +++ b/api/tests/unit_tests/domain/tracing/test_normalization.py @@ -8,7 +8,7 @@ from dify_trace_sdk.model import Cost, Span, SpanKind, SpanStatus, TokenUsage, Trace from pydantic import ValidationError -from services.tracing.normalization import ( +from domain.tracing.normalization import ( legacy_outer_trace, normalize_agent_span, order_spans, 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..376932d9e29b67 --- /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.model_dump(mode="json"), 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..8fe1b1d90b659c --- /dev/null +++ b/api/tests/unit_tests/domain/tracing/test_trace_run_agg.py @@ -0,0 +1,779 @@ +from collections.abc import Mapping +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", +) -> dict[str, object]: + """The JSON form of a 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)], + ) + ], + ).model_dump(mode="json") + + +def ingest(run: TraceRunAgg, payload: Mapping[str, object], *, 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) == (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 = UUID(str(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.model_dump(mode="json")) + 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.model_dump(mode="json")) + + +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/extensions/application_services/test_tracing.py b/api/tests/unit_tests/extensions/application_services/test_tracing.py index 63ec791f20a1e7..9e079027b2eed3 100644 --- a/api/tests/unit_tests/extensions/application_services/test_tracing.py +++ b/api/tests/unit_tests/extensions/application_services/test_tracing.py @@ -43,7 +43,11 @@ def test_agent_collection_shares_the_providers_and_schedules_exports_on_the_trac tracing = build_tracing_services(database_client=sqlite_session_factory) agent_traces = tracing.agent_traces - assert isinstance(agent_traces._traces, ApplicationTraceRepository) + # 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 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 da8bd5da3655d8..2cc696c336ba2a 100644 --- a/api/tests/unit_tests/repositories/test_application_trace_repository.py +++ b/api/tests/unit_tests/repositories/test_application_trace_repository.py @@ -1,23 +1,28 @@ import json -from collections.abc import Callable, Iterator -from dataclasses import dataclass +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 +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 WorkflowRunCompleted +from dify_trace_sdk.events import TraceEvent, WorkflowRunCompleted from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace -from dify_trace_sdk.provider import SendResult 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 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 @@ -25,11 +30,13 @@ from models.enums import ConversationFromSource from models.model import App, AppMode, Message, TraceAppConfig from models.workflow import WorkflowNodeExecutionModel, WorkflowRun -from repositories import application_trace_repository as repository_module from repositories.application_trace_repository import ApplicationTraceRepository 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 TraceInvocation, TraceRun, TraceRunState, TraceRunStatus +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, @@ -37,11 +44,17 @@ TracingStoreUnavailableError, ) from services.tracing.providers import ProviderConfigCodec -from services.tracing.records import ModerationTraceInfo, TraceMessage, WorkflowTraceInfo -from services.tracing.recovery import _workflow_event +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_app, make_message, make_workflow +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, @@ -54,22 +67,52 @@ "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, 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 self.repository.invocation( - run_id=self.run.id, agent_id=identifier("agent"), node_execution_id=None, attempt=0, call_key="first" - ) + return open_invocation(self, node_execution_id=None, call_key="first") def seed(sessions: sessionmaker[Session], name: str) -> Collection: @@ -81,14 +124,7 @@ def seed(sessions: sessionmaker[Session], name: str) -> Collection: 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": False, "capture_agent": True} - ), - ) + make_app(app_id=app_id, tenant_id=tenant_id, mode=AppMode.AGENT, tracing=json.dumps(LANGFUSE_POLICY)) ) session.add( make_message( @@ -118,7 +154,8 @@ def seed(sessions: sessionmaker[Session], name: str) -> Collection: ) ) repository = ApplicationTraceRepository(sessions) - run = repository.start( + scheduled: list[str] = [] + run = collector(repository, trace_records(sessions), scheduled).start( tenant_id=tenant_id, app_id=app_id, source_kind="message", @@ -127,7 +164,15 @@ def seed(sessions: sessionmaker[Session], name: str) -> Collection: workflow_id=None, ) assert run is not None - return Collection(repository, sessions, run) + 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) 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: @@ -151,14 +196,21 @@ def batch_for(run: TraceRun, invocation: TraceInvocation) -> AgentTraceBatch: ) -def ingest(collection: Collection, batch: AgentTraceBatch) -> tuple[int, str]: - return collection.repository.ingest(batch, max_spans=100, max_bytes=1_000_000) +def ingest(collection: Collection, batch: AgentTraceBatch) -> int: + return collection.service().ingest(batch) -def terminal(collection: Collection) -> None: - run = collection.run +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) as locked: + assert locked is not None + locked.run.ingest(batch.model_dump(mode="json"), 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=run.root_span_id, + span_id=root_span_id, parent_span_id=None, name="Agent turn", kind=SpanKind.CHAIN, @@ -166,7 +218,7 @@ def terminal(collection: Collection) -> None: end_time=datetime.fromtimestamp(3, UTC), status=SpanStatus.OK, ) - trace = Trace( + return Trace( trace_id=run.trace_id, external_trace_id=None, session_id=run.session_id, @@ -174,11 +226,82 @@ def terminal(collection: Collection) -> None: parent=None, spans=(root,), ) - collection.repository.terminal(run.id, trace.model_dump(mode="json"), root.span_id) -def claim(collection: Collection) -> tuple[TraceRun, str] | None: - return collection.repository.claim(collection.run.id, collection_timeout=60, lease_seconds=30) +def capture_terminal(collection: Collection, trace: Trace) -> None: + with collection.repository.change(collection.run.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) 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: @@ -224,10 +347,9 @@ class StoredSource: kind: Literal["workflow", "node"] row: dict[str, object] requests: list[str] - scheduled: list[str] def start(self) -> TraceRun: - run = self.case.repository.start( + run = self.service().start( tenant_id=self.case.run.tenant_id, app_id=self.case.run.app_id, source_kind=self.kind, @@ -239,17 +361,7 @@ def start(self) -> TraceRun: return run def service(self) -> ApplicationTracingService: - return ApplicationTracingService( - traces=self.case.repository, - records=TraceRecordRepository(self.case.sessions, files_url="https://files.test", enterprise_details=False), - builder=CanonicalTraceBuilder(), - # No provider is installed, so a delivery stops before it reaches any destination. - providers=InstalledTraceProviders(), - codec=ProviderConfigCodec(WorkspaceBoundCipher()), - egress=RecordingEgress([], grpc_target=None), - schedule_export=self.scheduled.append, - settings=AGENT_TRACE_SETTINGS, - ) + return self.case.service() def status(self, *, node_execution_id: str | None) -> TraceRunStatus: return self.service().status( @@ -261,6 +373,19 @@ def status(self, *, node_execution_id: str | None) -> TraceRunStatus: ) +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"]) @@ -369,7 +494,7 @@ def execute_sql( event.listen(sqlite_engine, "commit", ended) event.listen(sqlite_engine, "rollback", ended) try: - yield StoredSource(case, kind, row, requests, scheduled=[]) + yield StoredSource(case, kind, row, requests) finally: event.remove(sqlite_engine, "begin", begun) event.remove(sqlite_engine, "commit", ended) @@ -398,7 +523,7 @@ def test_status_and_retry_use_logstore_without_sql_source_rows(stored: StoredSou node_execution_id=None, ) assert result.status is TraceRunState.READY - assert stored.scheduled == [run.id] + assert stored.case.scheduled == [run.id] assert len(stored.requests) == 1 @staticmethod @@ -429,8 +554,11 @@ def test_foreign_app_owner_is_rejected_before_logstore_request(stored: StoredSou 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, @@ -445,7 +573,7 @@ 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, @@ -462,25 +590,22 @@ 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=TraceRunState.ACCEPTED, - error_code=None, - result=SendResult( - trace_id="provider-id", trace_url=None, span_urls={invocation.span_id: "https://trace.test/span"} - ), - retry_delay=None, - ) + token = claim(case) + assert token is not None + with case.repository.change(run.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")) @@ -489,16 +614,20 @@ def test_workflow_node_deep_link_is_scoped_to_durable_invocation(stored: StoredS def test_preview_settlement_and_workflow_recovery_use_configured_source(stored: StoredSource) -> None: run = stored.start() if stored.kind == "workflow": - event = _workflow_event(stored.case.repository, run, naive_utc_now()) - assert event is not None - assert event.payload == WorkflowRunCompleted( + records = AssembledEvents(trace_records(stored.case.sessions)) + service = collector(stored.case.repository, records, stored.case.scheduled) + stored.row["status"] = "running" + 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 (event.tenant_id, event.app_id) == (run.tenant_id, run.app_id) - assert event.user_id == stored.row["created_by"] - assert event.correlation.session_id == run.session_id - stored.row["status"] = "running" - assert _workflow_event(stored.case.repository, run, naive_utc_now()) is 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: # Delivery settles a finished preview from its node execution before claiming the export. stored.service().deliver(run.id) @@ -641,7 +770,7 @@ def test_sql_workflow_source_still_uses_scoped_configured_reader( def race_before_insert(monkeypatch: pytest.MonkeyPatch, competitor: Callable[[], None]) -> None: - """Run ``competitor`` after the caller's existence check, when it allocates its new run's identifiers.""" + """Run ``competitor`` after the caller's existence checks, when it allocates its new run's identifiers.""" raced: list[bool] = [] def allocate() -> UUID: @@ -650,7 +779,7 @@ def allocate() -> UUID: competitor() return uuidv7() - monkeypatch.setattr(repository_module, "uuidv7", allocate) + monkeypatch.setattr(service_module, "uuidv7", allocate) def drop_run(case: Collection) -> None: @@ -668,7 +797,7 @@ def run_count(case: Collection) -> int | None: def start_message(case: Collection, session_id: str) -> TraceRun | None: - return case.repository.start( + return case.service().start( tenant_id=case.run.tenant_id, app_id=case.run.app_id, source_kind="message", @@ -678,6 +807,17 @@ def start_message(case: Collection, session_id: str) -> TraceRun | 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) -> Generator[LockedTraceRun | None]: + self.purge(retention_days=0, limit=100) + with super().change(run_id) as locked: + yield locked + + class TestApplicationTraceRepositoryBoundariesCases: @staticmethod def test_concurrent_start_returns_the_first_frozen_snapshot( @@ -701,9 +841,12 @@ def test_initialization_failure_yields_to_a_concurrently_frozen_start( ) -> 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"))) - current = case.repository.record_initialization_failure( + # 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 @@ -721,7 +864,7 @@ def test_identifier_collision_is_not_mistaken_for_a_concurrent_start( case = seed(sqlite_session_factory, "colliding-start") other = seed(sqlite_session_factory, "colliding-owner") drop_run(case) - monkeypatch.setattr(repository_module, "uuidv7", lambda: UUID(other.run.id)) + 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) @@ -739,7 +882,7 @@ 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", @@ -747,6 +890,10 @@ def test_malformed_policy_freezes_disabled_decision( 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 is TraceRunState.NOT_COLLECTED assert run.provider_config == {} @@ -757,42 +904,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 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"]) @@ -834,7 +987,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 @@ -848,7 +1002,7 @@ def test_first_terminal_is_immutable_and_revoked_owner_is_never_exported( assert failed is not None assert failed.error_code == "tracing_owner_not_found" assert failed.state is TraceRunState.FAILED - case.repository.terminal(case.run.id, {"replacement": True}, "f" * 16) + capture_terminal(case, replacement) frozen = case.repository.get(case.run.id) assert frozen is not None assert frozen.outer_trace == before.outer_trace @@ -875,7 +1029,14 @@ def test_retry_rejects_ineligible_states( assert row is not None row.state = state with pytest.raises(error, match=code): - case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) + 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 @@ -885,17 +1046,34 @@ def test_retry_unknown_run_and_invocation_reject_before_creation( sqlite_session_factory: sessionmaker[Session], ) -> None: case = seed(sqlite_session_factory, "retry-unknown") + purged = replace(case, repository=PurgedBeforeLock(case.sessions)).service() with pytest.raises(TraceCollectionNotFoundError, match="tracing_run_not_found"): - case.repository.retry(run_id=identifier("missing"), acknowledge_possible_duplicates=True) + 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"): - case.repository.invocation( - run_id=identifier("missing"), + 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: @@ -921,8 +1099,511 @@ def test_retention_removes_all_payload_rows_and_preserves_newer_runs( row.created_at = naive_utc_now() - timedelta(days=8) 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 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) 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) 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.model_dump(mode="json"), 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) 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.model_dump(mode="json"), 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) 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) 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) 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.model_dump(mode="json"), 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) 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_a_missing_run(sqlite_session_factory: sessionmaker[Session]) -> None: + case = seed(sqlite_session_factory, "missing") + with case.repository.change(identifier("missing-run")) 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) 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.model_dump(mode="json"), 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/services/tracing/test_provider_configs.py b/api/tests/unit_tests/services/tracing/test_provider_configs.py index 6ba64cc23ef4cf..b1922e92483a73 100644 --- a/api/tests/unit_tests/services/tracing/test_provider_configs.py +++ b/api/tests/unit_tests/services/tracing/test_provider_configs.py @@ -6,12 +6,8 @@ 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, - TraceProviderSupport, -) -from services.tracing.providers import ProviderConfigCodec, configured_secrets +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" diff --git a/api/tests/unit_tests/services/tracing/test_recovery.py b/api/tests/unit_tests/services/tracing/test_recovery.py index 301d70d8692181..40c5225030c0dd 100644 --- a/api/tests/unit_tests/services/tracing/test_recovery.py +++ b/api/tests/unit_tests/services/tracing/test_recovery.py @@ -49,6 +49,7 @@ ) 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 @@ -65,7 +66,7 @@ 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, TraceInvocation, TraceRun, TraceRunState +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 @@ -104,9 +105,11 @@ class Collection: run: TraceRun def invocation(self, name: str) -> TraceInvocation: - return self.repository.invocation( - run_id=self.run.id, agent_id=identifier("agent"), node_execution_id=None, attempt=0, call_key=name - ) + with self.repository.change(self.run.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]): @@ -142,6 +145,32 @@ def send(self, config: LangfuseConfig, trace: Trace) -> SendResult: ) +def record_source(sessions: sessionmaker[Session]) -> TraceRecordRepository: + return TraceRecordRepository(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}") @@ -190,7 +219,9 @@ def seed(sessions: sessionmaker[Session], name: str, *, capture_content: bool, c ) ) repository = ApplicationTraceRepository(sessions) - run = repository.start( + 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", @@ -269,8 +300,45 @@ def batch_for(run: TraceRun, invocation: TraceInvocation, *, final: bool, span_i ) -def ingest(collection: Collection, batch: AgentTraceBatch) -> tuple[int, str]: - return collection.repository.ingest(batch, max_spans=100, max_bytes=1_000_000) +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)) as locked: + assert locked is not None + locked.run.ingest(batch.model_dump(mode="json"), 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) 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]: @@ -289,13 +357,8 @@ class RecoveryCase: def recover(self, run: TraceRun, *, now: datetime) -> bool: """One reconciler recovery attempt of ``run`` at ``now``.""" - return recover_terminal( - run, - traces=self.collection.repository, - records=self.records, - collector=self.service, - collection_timeout=self.settings.collection_timeout_seconds, - now=now, + return recover( + self.collection, run, records=self.records, collector=self.service, settings=self.settings, now=now ) @@ -303,18 +366,9 @@ 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 = TraceRecordRepository(collection.sessions, files_url="https://files.test", enterprise_details=False) + records = record_source(collection.sessions) scheduled: list[str] = [] - service = ApplicationTracingService( - traces=collection.repository, - records=records, - builder=CanonicalTraceBuilder(), - providers=InstalledTraceProviders(ScriptedProvider), - codec=ProviderConfigCodec(WorkspaceBoundCipher()), - egress=destination, - schedule_export=scheduled.append, - settings=settings, - ) + service = tracing_service(collection.repository, records, destination, scheduled, settings) return RecoveryCase(collection, records, service, settings, now, scheduled, destination) @@ -613,7 +667,7 @@ def test_stopped_consumer_cannot_freeze_a_terminal_before_the_runtime_worker( assert stored.outer_trace is None, "a stopped SSE save is not the runtime terminal" assert queued == [] assert stopped.recover(stored, now=case.now) is False - assert case.collection.repository.claim(stored.id, collection_timeout=60, lease_seconds=30) is None + assert claim(case.collection, stored.id) is None service.deliver(stored.id) assert destination.requests == [] @@ -810,7 +864,7 @@ def unexpected_cancel(*_args: object, **_kwargs: object) -> None: assert root.status.value == ("ok" if runtime_outcome == "succeeded" else "error") assert len(store.saved) == (0 if snapshot_save_fails else 1) 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 case.destination.requests == [] ingest(case.collection, batch_for(run, pending_invocation, final=True, span_id="b" * 16)) @@ -848,7 +902,7 @@ def test_missing_runtime_stop_terminal_converges_without_inventing_success( "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]) @@ -1200,13 +1254,9 @@ def test_cross_tenant_source_is_rejected_before_the_record_repository_runs(case: assert app is not None app.tenant_id = identifier("other-tenant") records = _FailingRecords(AssertionError("a foreign source must never be assembled")) - recovered = recover_terminal( - case.collection.run, - traces=case.collection.repository, - records=records, - collector=case.service, - collection_timeout=case.settings.collection_timeout_seconds, - now=case.now, + 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 == [] @@ -1222,13 +1272,9 @@ def test_source_not_ready_for_assembly_is_retried_later(case: RecoveryCase) -> N message.message_metadata = complete_metadata() message.updated_at = case.now - timedelta(seconds=11) records = _FailingRecords(SourceNotReadyError()) - recovered = recover_terminal( - case.collection.run, - traces=case.collection.repository, - records=records, - collector=case.service, - collection_timeout=case.settings.collection_timeout_seconds, - now=case.now, + 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] == [ @@ -1293,7 +1339,7 @@ def add_workflow_run(case: RecoveryCase, status: WorkflowExecutionStatus, *, fin def start_workflow(case: RecoveryCase, workflow_run_id: str) -> TraceRun: - run = case.collection.repository.start( + run = case.service.start( tenant_id=case.collection.run.tenant_id, app_id=case.collection.run.app_id, source_kind="workflow", @@ -1499,4 +1545,4 @@ def test_worker_handoff_leaves_a_not_collected_run_untouched(case: RecoveryCase) assert stored is not None assert stored.outer_trace is None assert case.scheduled == [] - assert case.collection.repository.claim(case.collection.run.id, collection_timeout=60, lease_seconds=30) is None + assert claim(case.collection, case.collection.run.id) is None 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 72853cef81ef10..ccc648d5d37fee 100644 --- a/api/tests/unit_tests/services/tracing/test_runtime_bridge.py +++ b/api/tests/unit_tests/services/tracing/test_runtime_bridge.py @@ -33,9 +33,11 @@ from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom, WorkflowAppGenerateEntity from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeBuildContext +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, @@ -62,17 +64,9 @@ 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, - ApplicationTraceStore, - TraceRun, - TraceRunState, - TraceSourceKind, - workflow_has_agent_v2, -) +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.normalization import remap_outer_trace from services.tracing.providers import ProviderConfigCodec from services.tracing.records import WorkflowTraceInfo from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks @@ -132,23 +126,14 @@ def schedule(self, job: ExportJob, *, delay_seconds: int) -> None: class _FreezeFails(ApplicationTraceRepository): - """The real ledger, except that freezing a run fails with ``failure``.""" + """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 start( - self, - *, - tenant_id: str, - app_id: str, - source_kind: TraceSourceKind, - source_id: str, - session_id: str, - workflow_id: str | None, - ) -> TraceRun | None: + def app_tracing(self, *, tenant_id: str, app_id: str) -> AppTracingSnapshot | None: raise self._failure @@ -158,14 +143,16 @@ def record_source(sessions: sessionmaker[Session]) -> TraceRecordRepository: def tracing_service( sessions: sessionmaker[Session], - traces: ApplicationTraceStore, + 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( - traces=traces, + runs=traces, + run_query=traces, + sources=traces, records=record_source(sessions), builder=CanonicalTraceBuilder(), providers=InstalledTraceProviders(LangfuseProvider), @@ -272,7 +259,7 @@ def seed(sessions: sessionmaker[Session], name: str) -> Collection: ) ) repository = ApplicationTraceRepository(sessions) - run = repository.start( + run = tracing_service(sessions, repository, no_egress(), [], AGENT_TRACE_SETTINGS).start( tenant_id=tenant_id, app_id=app_id, source_kind="message", @@ -330,6 +317,20 @@ def trace_event(run: TraceRun, payload: TracePayload) -> TraceEvent: ) +def capture(traces: ApplicationTraceRepository, run_id: str, trace: Trace) -> None: + """Keep a business terminal in its run without scheduling the run's export.""" + with traces.change(run_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.""" + with traces.change(str(batch.run_id)) as locked: + assert locked is not None + locked.run.ingest(batch.model_dump(mode="json"), 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 @@ -350,7 +351,7 @@ def terminal(collection: Collection) -> None: parent=None, spans=(root,), ) - collection.repository.terminal(run.id, trace.model_dump(mode="json"), root.span_id) + capture(collection.repository, run.id, trace) @dataclass(frozen=True) @@ -360,7 +361,7 @@ class Case: scheduled: list[str] -def handoff(collection: Collection, traces: ApplicationTraceStore, egress: Egress) -> Case: +def handoff(collection: Collection, traces: ApplicationTraceRepository, egress: Egress) -> Case: """The Agent collector over ``traces``, exporting through ``egress``.""" scheduled: list[str] = [] return Case( @@ -462,7 +463,7 @@ def workflow(case: Case, name: str, *, has_agent: bool, snapshot: bool) -> tuple 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", @@ -768,13 +769,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) 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, @@ -807,7 +806,8 @@ 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 ), @@ -848,7 +848,7 @@ def run(self, *, app_id: str, source_kind: TraceSourceKind, source_id: str) -> T def bridge_hooks( - sessions: sessionmaker[Session], traces: ApplicationTraceStore, settings: AgentTraceSettings + sessions: sessionmaker[Session], traces: ApplicationTraceRepository, settings: AgentTraceSettings ) -> RuntimeApplicationTracingHooks: return RuntimeApplicationTracingHooks(tracing_service(sessions, traces, no_egress(), [], settings)) @@ -1085,7 +1085,7 @@ def test_graph_string_with_agent_is_recognized() -> None: def test_worker_terminal_handoff_preserves_outcome_without_sse( bridge: Bridge, outcome: Literal["succeeded", "failed", "cancelled"] ) -> None: - started = bridge.traces.start( + started = tracing_service(bridge.sessions, bridge.traces, no_egress(), [], AGENT_TRACE_SETTINGS).start( tenant_id="tenant", app_id="agent-app", source_kind="message", @@ -1420,10 +1420,10 @@ def test_hidden_child_skip_keeps_normal_root_agent_trace_assemblable( ) outer = CanonicalTraceBuilder().build(info) assert outer is not None - outer = remap_outer_trace(outer, trace_id=run.trace_id, capture_content=False, synthetic_span_ids=frozenset()) - traces.terminal(run.id, outer.model_dump(mode="json"), outer.root.span_id) + capture(traces, run.id, outer) [invocation] = traces.invocations(run.id) - traces.ingest( + ingest( + traces, AgentTraceBatch( run_id=UUID(run.id), invocation_id=UUID(invocation.id), diff --git a/api/tests/unit_tests/services/tracing/test_service.py b/api/tests/unit_tests/services/tracing/test_service.py index d113757b5a270a..1dfe6b4a4a3276 100644 --- a/api/tests/unit_tests/services/tracing/test_service.py +++ b/api/tests/unit_tests/services/tracing/test_service.py @@ -1,6 +1,7 @@ import json import threading -from collections.abc import Callable, Iterator, Mapping, Sequence +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 @@ -23,6 +24,8 @@ 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, @@ -36,20 +39,14 @@ from models.enums import ConversationFromSource, CreatorUserRole from models.model import App, AppMode, Message, TraceAppConfig from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun -from repositories import application_trace_repository as repository_module from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository from repositories.application_trace_repository import ApplicationTraceRepository 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.canonical import CanonicalTraceBuilder -from services.tracing.collection import ( - AgentTraceSettings, - TraceInvocation, - TraceRun, - TraceRunState, - TraceRunStatus, - TraceSourceKind, -) +from services.tracing.collection import AgentTraceSettings, TraceRun, TraceRunStatus +from services.tracing.collection_ports import LockedTraceRun from services.tracing.errors import ( TraceCollectionConflictError, TraceCollectionLimitError, @@ -85,15 +82,29 @@ class _ApplicationTraceCollectionCollection: sessions: sessionmaker[Session] 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) as locked: + assert locked is not None + yield locked.run + def invocation(self, name: str = "first") -> TraceInvocation: - return self.repository.invocation( - run_id=self.run.id, + 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.""" @@ -189,7 +200,7 @@ def trace_records(sessions: sessionmaker[Session]) -> TraceRecordRepository: @dataclass class ServicePorts: - """The service's collaborators besides its ledger; a test replaces those it constrains. + """The service's collaborators besides its repository; a test replaces those it constrains. ``scheduled`` lists the runs the service scheduled for export. """ @@ -201,9 +212,17 @@ class ServicePorts: 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( - traces=case.repository, - records=trace_records(case.sessions), + runs=repository, + run_query=repository, + sources=repository, + records=trace_records(sessions), builder=CanonicalTraceBuilder(), providers=self.providers, codec=ProviderConfigCodec(self.cipher), @@ -284,7 +303,7 @@ def _application_trace_collection_seed( ) ) repository = ApplicationTraceRepository(sessions) - run = repository.start( + ServicePorts(_application_trace_collection_no_egress()).service_on(repository, sessions).start( tenant_id=tenant_id, app_id=app_id, source_kind="message", @@ -292,6 +311,8 @@ def _application_trace_collection_seed( 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) @@ -383,8 +404,10 @@ 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: @@ -406,13 +429,39 @@ def _application_trace_collection_terminal(collection: _ApplicationTraceCollecti parent=None, spans=(root,), ) - collection.repository.terminal(run.id, trace.model_dump(mode="json"), root.span_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) 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[TraceRun, 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: @@ -427,20 +476,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=TraceRunState.ACCEPTED, - error_code=None, - result=SendResult(trace_id="provider-id", trace_url=None, span_urls={}), - retry_delay=None, - ) - 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 @@ -456,7 +498,7 @@ def test_conflicting_batch_id_rolls_back_without_replacing_payload( changed = batch.model_copy(update={"final": True}) with pytest.raises(TraceCollectionConflictError, match="tracing_batch_conflict"): _application_trace_collection_ingest(collection, changed) - invocations, _ = collection.repository.collected(collection.run.id) + invocations = collection.repository.invocations(collection.run.id) assert invocations[0].finalized is False with collection.sessions() as session: assert session.scalar(select(ApplicationTraceBatch.final)) is False @@ -474,11 +516,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()) + 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"): - 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 + 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 @@ -489,8 +532,7 @@ def test_wrong_trace_id_is_rejected_before_persistence(collection: _ApplicationT altered = batch.model_copy(update={"spans": [batch.spans[0].model_copy(update={"trace_id": "f" * 32})]}) 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( @@ -518,7 +560,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" @@ -538,9 +580,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(TraceCollectionConflictError, 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, @@ -582,14 +623,7 @@ def test_expired_delivery_lease_accepts_confirmed_late_ack_for_same_attempt( current = collection.repository.get(collection.run.id) assert current is not None assert current.state is TraceRunState.UNCERTAIN - collection.repository.complete( - collection.run.id, - leased[1], - state=TraceRunState.ACCEPTED, - error_code=None, - result=SendResult(trace_id="late", trace_url=None, span_urls={}), - retry_delay=None, - ) + _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 @@ -624,7 +658,8 @@ 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", @@ -679,9 +714,11 @@ def test_delivery_checkpoint_controls_replays( @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, @@ -690,9 +727,16 @@ def test_source_owner_checks_app_and_tenant(collection: _ApplicationTraceCollect ) 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", @@ -702,6 +746,12 @@ def test_source_owner_checks_app_and_tenant(collection: _ApplicationTraceCollect ) 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( @@ -765,8 +815,7 @@ def test_capture_agent_disabled_preserves_outer_only_content_policy( 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( @@ -929,8 +978,7 @@ def test_status_uses_frozen_destination_and_selects_latest_node_attempt( assert config is not None 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, @@ -942,17 +990,14 @@ def test_status_uses_frozen_destination_and_selects_latest_node_attempt( _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=TraceRunState.ACCEPTED, - error_code=None, - result=SendResult( + SendResult( trace_id="provider-trace", trace_url="https://frozen.example/trace", span_urls={invocation.span_id: "https://frozen.example/span"}, ), - retry_delay=None, ) service = ServicePorts( _application_trace_service_boundaries_no_egress(), @@ -1001,7 +1046,8 @@ 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", @@ -1010,7 +1056,6 @@ def test_new_run_freezes_updated_project_without_changing_previous_run( workflow_id=None, ) assert new_run is not None - service = ServicePorts(_application_trace_service_boundaries_no_egress()).service(case) previous = service.status(**source(case.run), node_execution_id=None) assert previous.project_url == "https://langfuse.test/project/original" assert ( @@ -1109,14 +1154,8 @@ 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=TraceRunState.UNCERTAIN, - error_code="provider_acceptance_unknown", - result=None, - retry_delay=None, - ) + 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"): @@ -1198,14 +1237,15 @@ def test_unreadable_frozen_configuration_fails_before_provider_io( freeze_config(case, "langfuse", {**FROZEN_LANGFUSE, "host": "private-host-without-scheme"}) cipher = UnreadableSecretCipher() if unreadable == "decryption" else StoredSecretCipher() destination = RecordingEgress([], grpc_target=None) - ServicePorts(destination, cipher=cipher).service(case).deliver(case.run.id) + 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"): - case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) + service.retry(**source(case.run), acknowledge_possible_duplicates=True, node_execution_id=None) @staticmethod def test_unavailable_frozen_provider_fails_before_decryption( @@ -1267,7 +1307,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", @@ -1276,8 +1318,7 @@ def test_preview_terminal_uses_persisted_node_state_and_frozen_privacy( session_id="preview-session", ) assert run is not None - destination = _application_trace_service_boundaries_accepted("preview-trace") - ServicePorts(destination).service(case).deliver(run.id) + service.deliver(run.id) current = case.repository.get(run.id) assert current is not None if status == WorkflowNodeExecutionStatus.RUNNING: @@ -1437,8 +1478,11 @@ def test_invalid_tree_is_permanent_before_decryption_or_provider_io( assert current.state is TraceRunState.FAILED assert current.error_code == "tracing_payload_invalid" assert current.attempts == 1 + # 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"): - case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) + 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.""" @@ -1478,8 +1522,7 @@ def status(self, node_execution_id: str | None = None) -> TraceRunStatus: 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, + invocation = self.case.open_invocation( agent_id=_application_trace_node_alias_identifier("referenced-agent"), node_execution_id=str(self.node["id"]), attempt=attempt, @@ -1493,17 +1536,14 @@ def accept(self, *, spans: bool = True) -> tuple[str, str]: _application_trace_node_alias_terminal(self.case) lease = _application_trace_node_alias_claim(self.case) assert lease is not None - self.case.repository.complete( - self.case.run.id, + _accept_export( + self.case, lease[1], - state=TraceRunState.ACCEPTED, - error_code=None, - result=SendResult( + SendResult( trace_id="provider-trace", trace_url="https://provider.test/trace", span_urls=span_urls if spans else {}, ), - retry_delay=None, ) return "https://provider.test/trace", "https://provider.test/span/1" @@ -1627,7 +1667,8 @@ def execute_sql( event.listen(sqlite_engine, "commit", ended) event.listen(sqlite_engine, "rollback", ended) try: - run = case.repository.start( + 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", @@ -1679,13 +1720,10 @@ def test_alias_without_internal_invocation_returns_parent_status( _application_trace_node_alias_terminal(alias.case) lease = _application_trace_node_alias_claim(alias.case) assert lease is not None - alias.case.repository.complete( - alias.case.run.id, + _accept_export( + alias.case, lease[1], - state=TraceRunState.ACCEPTED, - error_code=None, - result=SendResult(trace_id="outer-only", trace_url="https://provider.test/outer", span_urls={}), - retry_delay=None, + 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) @@ -1697,7 +1735,7 @@ def test_alias_without_internal_invocation_returns_parent_status( @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.case.repository.start( + node_run = alias.service().start( tenant_id=alias.case.run.tenant_id, app_id=alias.case.run.app_id, source_kind="node", @@ -1719,8 +1757,7 @@ def test_standalone_node_snapshot_wins_over_parent(alias: NodeAlias, state: Trac @staticmethod def test_explicit_other_node_filter_cannot_select_sibling_span(alias: NodeAlias) -> None: - sibling = alias.case.repository.invocation( - run_id=alias.case.run.id, + 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, @@ -1861,8 +1898,18 @@ def request( ) +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.""" + """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) @@ -1871,36 +1918,49 @@ def __init__(self, sessions: sessionmaker[Session], *, transient_failure: bool) self.renewed = threading.Event() @override - def renew_export_lease(self, run_id: str, token: str) -> bool: - self.renewal_threads.append(threading.get_ident()) - if self._transient_failure and len(self.renewal_threads) == 1: - raise RuntimeError("private-provider-value") - renewal = super().renew_export_lease(run_id, token) - self.renewed.set() - return renewal + @contextmanager + def change(self, run_id: str) -> Generator[LockedTraceRun | None]: + renewal = False + with super().change(run_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 - def complete( - self, - run_id: str, - token: str, - *, - state: TraceRunState, - error_code: str | None, - result: SendResult | None, - retry_delay: int | None, - ) -> None: - raise TracingStoreUnavailableError("database unavailable") + @contextmanager + def change(self, run_id: str) -> Generator[LockedTraceRun | None]: + with super().change(run_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: @@ -1932,30 +1992,27 @@ def case(sqlite_session_factory: sessionmaker[Session]) -> _ExportLeaseCollectio return collection @staticmethod - def test_renewal_keeps_long_export_alive_then_missing_heartbeat_expires( - case: _ExportLeaseCollection, monkeypatch: pytest.MonkeyPatch - ) -> None: - clock = [naive_utc_now()] - monkeypatch.setattr(repository_module, "naive_utc_now", lambda: clock[0]) - leased = _export_lease_claim(case) + 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] - clock[0] += timedelta(seconds=20) - assert case.repository.renew_export_lease(case.run.id, token) - clock[0] += timedelta(seconds=20) - assert _export_lease_claim(case) is None + 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) < clock[0] - clock[0] += timedelta(seconds=11) - assert _export_lease_claim(case) is 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 case.repository.renew_export_lease(case.run.id, token) + assert not _renew_export_lease(case, token, now) @staticmethod def test_old_rows_without_heartbeat_use_started_at(case: _ExportLeaseCollection) -> None: @@ -1976,15 +2033,8 @@ def test_wrong_token_cannot_renew_or_complete(case: _ExportLeaseCollection) -> N leased = _export_lease_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=TraceRunState.ACCEPTED, - error_code=None, - result=result("wrong"), - retry_delay=None, - ) + 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 @@ -1999,9 +2049,7 @@ def test_late_success_wins_before_manual_retry_and_is_idempotent(case: _ExportLe trace_url="https://example.com/trace", span_urls={"a" * 16: "https://example.com/span"}, ) - case.repository.complete( - case.run.id, token, state=TraceRunState.ACCEPTED, error_code=None, result=confirmed, retry_delay=None - ) + _accept_export(case, token, confirmed) current = _stored(case) assert current.state == TraceRunState.ACCEPTED assert current.error_code is None @@ -2009,57 +2057,29 @@ def test_late_success_wins_before_manual_retry_and_is_idempotent(case: _ExportLe assert current.accepted_at is not None assert current.trace_url == confirmed.trace_url assert current.span_urls == {"a" * 16: "https://example.com/span"} - case.repository.complete( - case.run.id, - token, - state=TraceRunState.ACCEPTED, - error_code=None, - result=result("duplicate-ack"), - retry_delay=None, - ) + _accept_export(case, token, result("duplicate-ack")) assert _stored(case).provider_trace_id == "confirmed" - with pytest.raises(TraceCollectionConflictError, match="tracing_already_accepted"): - case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) + 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) - 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=TraceRunState.ACCEPTED, - error_code=None, - result=result("old"), - retry_delay=None, - ) + _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 - case.repository.complete( - case.run.id, - old_token, - state=TraceRunState.ACCEPTED, - error_code=None, - result=result("old"), - retry_delay=None, - ) + _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 - case.repository.complete( - case.run.id, - second[1], - state=TraceRunState.ACCEPTED, - error_code=None, - result=result("new"), - retry_delay=None, - ) + _accept_export(case, second[1], result("new")) assert _stored(case).provider_trace_id == "new" @staticmethod @@ -2068,9 +2088,14 @@ def test_late_non_success_cannot_overwrite_lease_expiry( case: _ExportLeaseCollection, late_state: TraceRunState ) -> None: token = _expire(case) - case.repository.complete( - case.run.id, token, state=late_state, error_code="late-failure", result=None, retry_delay=None - ) + 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" @@ -2082,9 +2107,7 @@ def test_other_uncertain_reason_cannot_be_promoted_by_late_ack(case: _ExportLeas 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=TraceRunState.ACCEPTED, error_code=None, result=result("late"), retry_delay=None - ) + _accept_export(case, token, result("late")) assert _stored(case).state == TraceRunState.UNCERTAIN @staticmethod @@ -2105,7 +2128,8 @@ def test_service_heartbeats_during_export_and_stops_on_every_outcome( reply: EgressResponse | Exception, ) -> None: clock = [naive_utc_now()] - monkeypatch.setattr(repository_module, "naive_utc_now", lambda: clock[0]) + # 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]] = [] @@ -2114,7 +2138,7 @@ def while_exporting() -> None: # transaction can prevent a second claimant from marking it uncertain. clock[0] += timedelta(seconds=3) heartbeat_committed = renewals.renewed.wait(3) - second_claim = case.repository.claim(case.run.id, collection_timeout=60, lease_seconds=1) + second_claim = _claim_at(case, lease_seconds=1, now=clock[0]) observed.append((heartbeat_committed, second_claim is None, _stored(case).state)) ports = ServicePorts( @@ -2131,8 +2155,11 @@ def while_exporting() -> None: @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(_export_lease_accepted("ack")).service(unavailable).deliver(case.run.id) + 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()) 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 192abb2428a994..d02af7d055d76e 100644 --- a/api/tests/unit_tests/tasks/test_application_trace_task.py +++ b/api/tests/unit_tests/tasks/test_application_trace_task.py @@ -11,6 +11,7 @@ from sqlalchemy.orm import Session, sessionmaker 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 @@ -19,7 +20,7 @@ from models.model import AppMode, Message from repositories.trace_record_repository import TraceRecordRepository from services.tracing.canonical import CanonicalTraceBuilder -from services.tracing.collection import AgentTraceSettings, TraceRunState +from services.tracing.collection import AgentTraceSettings from services.tracing.egress import SSRFEgress from services.tracing.providers import ProviderConfigCodec from services.tracing.records import MessageTraceInfo @@ -48,7 +49,9 @@ def bind_worker( ) -> None: """Hand the tasks a worker over the seeded ledger with the installed providers and the SSRF egress.""" worker = ApplicationTracingService( - traces=collection.repository, + runs=collection.repository, + run_query=collection.repository, + sources=collection.repository, records=TraceRecordRepository(collection.sessions, files_url="https://files.test", enterprise_details=False), builder=CanonicalTraceBuilder(), providers=EntryPointTraceProviderCatalog(), From 98bc515d01b2da7815c53235156abefde7ba3e26 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Thu, 29 Oct 2026 00:05:46 +0800 Subject: [PATCH 04/17] feat(tracing): nest workflow traces under the tool node that started them A workflow run started by a workflow tool node now joins the trace of the run that started it, hanging from that tool node, when both apps export to the same destination. Otherwise it keeps a trace of its own that names its outer run. Workflow tool nodes pass the chain of outer runs down to the runs they start, so a run nested at any depth joins the outermost trace it shares a destination with. Delivery reads the outer runs' committed records and applies a pure domain placement rule. An Agent-collected outer run lends its frozen trace once collection has started. Trace units can now hang from any span of their trace. LangSmith hangs them from the trace root, because its run order needs the full ancestry, and MLflow still sends them as traces of their own. --- api/core/app/apps/advanced_chat/app_runner.py | 6 + api/core/app/apps/workflow/app_runner.py | 6 + api/core/app/apps/workflow_app_runner.py | 15 ++ api/core/app/entities/app_invoke_entities.py | 9 + api/core/helper/trace_id_helper.py | 6 + api/core/tools/workflow_as_tool/tool.py | 12 +- api/core/workflow/node_runtime.py | 8 +- api/domain/tracing/nested_workflow_trace.py | 73 +++++++ api/domain/tracing/normalization.py | 6 +- api/domain/tracing/trace_run_agg.py | 2 + .../application_services/tracing.py | 8 +- api/extensions/tracing.py | 16 +- .../src/dify_trace_langfuse/provider.py | 5 +- .../src/dify_trace_langfuse/test_provider.py | 25 ++- .../src/dify_trace_langsmith/provider.py | 11 +- .../src/dify_trace_langsmith/test_provider.py | 20 ++ .../src/dify_trace_mlflow/_spans.py | 1 + .../src/dify_trace_mlflow/test_provider.py | 17 ++ .../src/dify_trace_opik/provider.py | 9 +- .../src/dify_trace_opik/test_provider.py | 16 ++ .../trace-sdk/src/dify_trace_sdk/events.py | 20 +- .../trace-sdk/src/dify_trace_sdk/model.py | 12 +- .../trace-sdk/src/dify_trace_sdk/testing.py | 48 +++++ .../src/dify_trace_sdk/test_emitter.py | 4 +- .../src/dify_trace_sdk/test_model.py | 28 +++ .../src/dify_trace_sdk/test_otlp.py | 14 +- .../src/dify_trace_sdk/test_testing.py | 9 +- .../src/dify_trace_weave/provider.py | 2 +- .../src/dify_trace_weave/test_provider.py | 17 ++ api/repositories/trace_record_repository.py | 45 +++- api/services/tracing/canonical.py | 14 +- api/services/tracing/delivery.py | 135 ++++++++++-- api/services/tracing/records.py | 3 + api/services/tracing/recovery.py | 3 +- api/services/tracing/scopes.py | 4 +- api/services/tracing/service.py | 37 ++-- .../core/agent/test_cot_agent_runner.py | 4 +- .../core/agent/test_fc_agent_runner.py | 4 +- .../apps/advanced_chat/test_app_generator.py | 8 +- .../test_generate_task_pipeline.py | 7 +- .../app/apps/agent_app/test_app_generator.py | 8 +- .../app/apps/agent_chat/test_app_generator.py | 4 +- .../core/app/apps/chat/test_app_generator.py | 4 +- .../app/apps/completion/test_app_generator.py | 4 +- .../core/app/apps/test_base_app_runner.py | 4 +- .../core/app/apps/test_workflow_app_runner.py | 14 +- .../app/apps/workflow/test_app_generator.py | 16 +- .../app/entities/test_app_invoke_entities.py | 4 +- ...st_easy_ui_based_generate_task_pipeline.py | 8 +- .../core/llm_generator/test_llm_generator.py | 3 +- .../core/moderation/test_input_moderation.py | 4 +- .../rag/retrieval/test_dataset_retrieval.py | 4 +- .../core/telemetry/test___init__.py | 4 +- .../unit_tests/core/telemetry/test_gateway.py | 4 +- .../unit_tests/core/tools/test_tool_engine.py | 4 +- .../core/tools/workflow_as_tool/test_tool.py | 38 +++- .../core/workflow/test_node_runtime.py | 40 +++- .../tracing/test_nested_workflow_trace.py | 108 ++++++++++ .../domain/tracing/test_normalization.py | 2 + .../domain/tracing/test_trace_run_agg.py | 1 + .../enterprise/telemetry/test_draft_trace.py | 4 +- .../telemetry/test_enterprise_trace.py | 1 + .../unit_tests/extensions/test_tracing.py | 35 +++- .../test_application_trace_repository.py | 33 +-- .../test_trace_record_repository.py | 143 ++++++++++++- .../services/test_workflow_persistence.py | 15 +- .../services/tracing/test_adapters.py | 5 +- .../services/tracing/test_canonical.py | 5 +- .../services/tracing/test_delivery.py | 193 +++++++++++++++++- .../services/tracing/test_recovery.py | 8 +- .../services/tracing/test_runtime_bridge.py | 126 +++++++++--- .../services/tracing/test_scopes.py | 4 +- .../services/tracing/test_service.py | 5 +- .../tasks/test_application_trace_task.py | 4 +- .../tasks/test_trace_delivery_task.py | 26 ++- 75 files changed, 1287 insertions(+), 287 deletions(-) create mode 100644 api/domain/tracing/nested_workflow_trace.py create mode 100644 api/tests/unit_tests/domain/tracing/test_nested_workflow_trace.py diff --git a/api/core/app/apps/advanced_chat/app_runner.py b/api/core/app/apps/advanced_chat/app_runner.py index 12687c89ee037b..5897360bb2fb2d 100644 --- a/api/core/app/apps/advanced_chat/app_runner.py +++ b/api/core/app/apps/advanced_chat/app_runner.py @@ -132,6 +132,8 @@ def prepare(self) -> PreparedWorkflowRun | None: user_from=user_from, app_type=get_credit_usage_app_type(app_config.app_mode), trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), + external_trace_id=self.application_generate_entity.extras.get("external_trace_id"), + parent_trace_context=self.application_generate_entity.extras.get("parent_trace_context"), call_depth=self.application_generate_entity.call_depth, ) elif self.application_generate_entity.single_iteration_run or self.application_generate_entity.single_loop_run: @@ -143,6 +145,8 @@ def prepare(self) -> PreparedWorkflowRun | None: user_id=self.application_generate_entity.user_id, app_type=get_credit_usage_app_type(app_config.app_mode), trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), + external_trace_id=self.application_generate_entity.extras.get("external_trace_id"), + parent_trace_context=self.application_generate_entity.extras.get("parent_trace_context"), ) else: inputs = self.application_generate_entity.inputs @@ -216,6 +220,8 @@ def prepare(self) -> PreparedWorkflowRun | None: root_node_id=root_node_id, app_type=get_credit_usage_app_type(app_config.app_mode), trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), + external_trace_id=self.application_generate_entity.extras.get("external_trace_id"), + parent_trace_context=self.application_generate_entity.extras.get("parent_trace_context"), call_depth=self.application_generate_entity.call_depth, ) diff --git a/api/core/app/apps/workflow/app_runner.py b/api/core/app/apps/workflow/app_runner.py index b2390ea19b156c..5171c6df4adb8d 100644 --- a/api/core/app/apps/workflow/app_runner.py +++ b/api/core/app/apps/workflow/app_runner.py @@ -97,6 +97,8 @@ def prepare(self) -> PreparedWorkflowRun | None: root_node_id=self._root_node_id, app_type=get_credit_usage_app_type(app_config.app_mode), trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), + external_trace_id=self.application_generate_entity.extras.get("external_trace_id"), + parent_trace_context=self.application_generate_entity.extras.get("parent_trace_context"), call_depth=self.application_generate_entity.call_depth, ) elif self.application_generate_entity.single_iteration_run or self.application_generate_entity.single_loop_run: @@ -107,6 +109,8 @@ def prepare(self) -> PreparedWorkflowRun | None: user_id=self.application_generate_entity.user_id, app_type=get_credit_usage_app_type(app_config.app_mode), trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), + external_trace_id=self.application_generate_entity.extras.get("external_trace_id"), + parent_trace_context=self.application_generate_entity.extras.get("parent_trace_context"), ) else: inputs = self.application_generate_entity.inputs @@ -153,6 +157,8 @@ def prepare(self) -> PreparedWorkflowRun | None: root_node_id=root_node_id, app_type=get_credit_usage_app_type(app_config.app_mode), trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), + external_trace_id=self.application_generate_entity.extras.get("external_trace_id"), + parent_trace_context=self.application_generate_entity.extras.get("parent_trace_context"), call_depth=self.application_generate_entity.call_depth, ) diff --git a/api/core/app/apps/workflow_app_runner.py b/api/core/app/apps/workflow_app_runner.py index 3df06d4b0df484..7fa0cef0dd683e 100644 --- a/api/core/app/apps/workflow_app_runner.py +++ b/api/core/app/apps/workflow_app_runner.py @@ -44,6 +44,7 @@ ) from core.app.layers.pause_state_persist_layer import PauseStateLayerConfig from core.credit_usage import CreditUsageAppType +from core.helper.trace_id_helper import ParentTraceContext from core.rag.entities import RetrievalSourceMetadata from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from core.repositories.human_input_repository import HumanInputFormRecord @@ -239,6 +240,8 @@ def _init_graph( root_node_id: str | None = None, app_type: CreditUsageAppType | None = None, trace_session_id: str | None = None, + external_trace_id: str | None = None, + parent_trace_context: ParentTraceContext | None = None, ) -> Graph: """ Init graph @@ -254,6 +257,8 @@ def _init_graph( invoke_from=invoke_from, app_type=app_type, trace_session_id=trace_session_id, + external_trace_id=external_trace_id, + parent_trace_context=parent_trace_context, ) graph_init_context = DifyGraphInitContext( execution_driver=self._execution_driver, @@ -290,6 +295,8 @@ def _prepare_single_node_execution( user_id: str, app_type: CreditUsageAppType | None = None, trace_session_id: str | None = None, + external_trace_id: str | None = None, + parent_trace_context: ParentTraceContext | None = None, ) -> tuple[Graph, VariablePool, RuntimeState]: """ Prepare graph, variable pool, and runtime state for single node execution @@ -328,6 +335,8 @@ def _prepare_single_node_execution( user_id=user_id, app_type=app_type, trace_session_id=trace_session_id, + external_trace_id=external_trace_id, + parent_trace_context=parent_trace_context, ) elif single_loop_run: graph, variable_pool = self._get_graph_and_variable_pool_for_single_node_run( @@ -339,6 +348,8 @@ def _prepare_single_node_execution( user_id=user_id, app_type=app_type, trace_session_id=trace_session_id, + external_trace_id=external_trace_id, + parent_trace_context=parent_trace_context, ) else: raise ValueError("Neither single_iteration_run nor single_loop_run is specified") @@ -358,6 +369,8 @@ def _get_graph_and_variable_pool_for_single_node_run( user_id: str = "", app_type: CreditUsageAppType | None = None, trace_session_id: str | None = None, + external_trace_id: str | None = None, + parent_trace_context: ParentTraceContext | None = None, ) -> tuple[Graph, VariablePool]: """ Get graph and variable pool for single node execution (iteration or loop). @@ -429,6 +442,8 @@ def _get_graph_and_variable_pool_for_single_node_run( invoke_from=InvokeFrom.DEBUGGER, app_type=app_type, trace_session_id=trace_session_id, + external_trace_id=external_trace_id, + parent_trace_context=parent_trace_context, ) graph_init_context = DifyGraphInitContext( execution_driver=self._execution_driver, diff --git a/api/core/app/entities/app_invoke_entities.py b/api/core/app/entities/app_invoke_entities.py index c2214eb088aabf..98cb29ff9c7820 100644 --- a/api/core/app/entities/app_invoke_entities.py +++ b/api/core/app/entities/app_invoke_entities.py @@ -17,6 +17,7 @@ normalize_credit_usage_created_by, ) from core.entities.provider_configuration import ProviderModelBundle +from core.helper.trace_id_helper import ParentTraceContext from graphon.file import File, FileUploadConfig from graphon.model_runtime.entities.model_entities import AIModelEntity from models.model import AppMode @@ -102,6 +103,10 @@ class DifyRunContext(BaseModel): app_type: CreditUsageAppType | None = None created_by: CreditUsageCreatedBy | None = None trace_session_id: str | None = None + external_trace_id: str | None = None + """The trace ID the run's caller supplied.""" + parent_trace_context: ParentTraceContext | None = None + """The workflow tool node that started this run; ``None`` unless the run is nested.""" workflow_tool_invocation_id: str | None = None @field_validator("created_by", mode="before") @@ -129,6 +134,8 @@ def build_dify_run_context( app_type: CreditUsageAppTypeInput = None, created_by: CreditUsageCreatedByInput = None, trace_session_id: str | None = None, + external_trace_id: str | None = None, + parent_trace_context: ParentTraceContext | None = None, extra_context: Mapping[str, Any] | None = None, ) -> dict[str, Any]: """ @@ -147,6 +154,8 @@ def build_dify_run_context( app_type=normalize_credit_usage_app_type(app_type) if app_type is not None else None, created_by=normalize_credit_usage_created_by(created_by) if created_by is not None else None, trace_session_id=trace_session_id, + external_trace_id=external_trace_id, + parent_trace_context=parent_trace_context, ) return run_context diff --git a/api/core/helper/trace_id_helper.py b/api/core/helper/trace_id_helper.py index e1ebd45e07478d..41e4686909a058 100644 --- a/api/core/helper/trace_id_helper.py +++ b/api/core/helper/trace_id_helper.py @@ -12,6 +12,12 @@ class ParentTraceContext(BaseModel): parent_workflow_run_id: StrictStr parent_node_execution_id: StrictStr | None = None + parent_app_id: StrictStr | None = None + """The app of the outer run; ``None`` in contexts saved before it was recorded.""" + parent_external_trace_id: StrictStr | None = None + """The trace ID the outer run's caller supplied; only an outermost run has one.""" + parent_trace_context: "ParentTraceContext | None" = None + """The outer run's own parent context; ``None`` when the outer run is not nested.""" model_config = ConfigDict(extra="forbid") diff --git a/api/core/tools/workflow_as_tool/tool.py b/api/core/tools/workflow_as_tool/tool.py index 9598b4ebb6494a..a9bf3ffb7a2bca 100644 --- a/api/core/tools/workflow_as_tool/tool.py +++ b/api/core/tools/workflow_as_tool/tool.py @@ -255,17 +255,9 @@ def fork_tool_runtime(self, runtime: ToolRuntime) -> WorkflowTool: forked._trace_session_id = self._trace_session_id return forked - def set_parent_trace_context( - self, - *, - parent_workflow_run_id: str, - parent_node_execution_id: str, - ) -> None: + def set_parent_trace_context(self, context: ParentTraceContext) -> None: """Attach outer workflow trace context without exposing it as tool input.""" - self._parent_trace_context = ParentTraceContext( - parent_workflow_run_id=parent_workflow_run_id, - parent_node_execution_id=parent_node_execution_id, - ) + self._parent_trace_context = context def clear_parent_trace_context(self) -> None: """Remove parent trace context before invoking this tool outside a nested workflow.""" diff --git a/api/core/workflow/node_runtime.py b/api/core/workflow/node_runtime.py index cd08dab797c9d4..79dbebea35aceb 100644 --- a/api/core/workflow/node_runtime.py +++ b/api/core/workflow/node_runtime.py @@ -610,6 +610,9 @@ def get_runtime( parent_trace_context = ParentTraceContext( parent_workflow_run_id=outer_workflow_run_id, parent_node_execution_id=node_execution_id, + parent_app_id=self._run_context.app_id, + parent_external_trace_id=self._run_context.external_trace_id, + parent_trace_context=self._run_context.parent_trace_context, ) if isinstance(self._run_context.trace_session_id, str) and self._run_context.trace_session_id: trace_session_id = self._run_context.trace_session_id @@ -647,10 +650,7 @@ def invoke( tool = runtime_binding.tool callback = DifyWorkflowCallbackHandler() if runtime_binding.parent_trace_context and hasattr(tool, "set_parent_trace_context"): - tool.set_parent_trace_context( - parent_workflow_run_id=runtime_binding.parent_trace_context.parent_workflow_run_id, - parent_node_execution_id=runtime_binding.parent_trace_context.parent_node_execution_id, - ) + tool.set_parent_trace_context(runtime_binding.parent_trace_context) elif hasattr(tool, "clear_parent_trace_context"): tool.clear_parent_trace_context() if runtime_binding.trace_session_id and hasattr(tool, "set_trace_session_id"): diff --git a/api/domain/tracing/nested_workflow_trace.py b/api/domain/tracing/nested_workflow_trace.py new file mode 100644 index 00000000000000..b7788cab66b061 --- /dev/null +++ b/api/domain/tracing/nested_workflow_trace.py @@ -0,0 +1,73 @@ +"""Place the trace of a nested workflow run in the trace of the run that started it. + +A nested run joins the trace of the run whose tool node started it when both +export to the same destination, and its root hangs from that tool node. That +outer run joined the trace above it by the same rule, so the joined trace is the +outermost one reachable through runs sharing the destination. A run whose spans +are collected for Agent tracing is always the root of its own trace. Every +identifier is derived from committed records, so the units need no coordination. +""" + +from collections.abc import Iterable +from dataclasses import dataclass + +from dify_trace_sdk.model import SpanRef, Trace + + +@dataclass(frozen=True, slots=True) +class OuterTrace: + """The trace an outer run's spans go to when that run is the root of its trace.""" + + trace_id: str + external_trace_id: str | None + """The caller-supplied trace ID the trace's root unit carries.""" + root: SpanRef + same_destination: bool + """Whether the nested run exports to the destination this trace goes to.""" + collected: bool + """Whether the run's spans are collected for Agent tracing, which keeps it the root of its trace.""" + + +@dataclass(frozen=True, slots=True) +class TracePlacement: + """Where the trace of a nested run hangs in an outer trace.""" + + trace_id: str + external_trace_id: str | None + parent: SpanRef + trace_root: SpanRef + + def place(self, trace: Trace) -> Trace: + """``trace`` as a unit of the outer trace, its root hanging from ``parent``.""" + root, *spans = trace.spans + return Trace( + trace_id=self.trace_id, + external_trace_id=self.external_trace_id, + session_id=trace.session_id, + user_id=trace.user_id, + parent=self.parent, + trace_root=self.trace_root, + spans=(root.model_copy(update={"parent_span_id": self.parent.span_id}), *spans), + ) + + +def nested_trace_placement(outer_traces: Iterable[OuterTrace], *, tool_node: SpanRef) -> TracePlacement | None: + """Where the trace of a nested run hangs; ``None`` when it starts a trace of its own. + + ``outer_traces`` follows the runs that started the nested run, innermost + first, and ends at the first run whose trace is not known yet. It is consumed + only as far as the decision needs. ``tool_node`` is the span of the tool node + that started the nested run. + """ + joined: OuterTrace | None = None + for outer in outer_traces: + if not outer.same_destination: + break + joined = outer + if outer.collected: + break + if joined is None: + return None + return TracePlacement( + trace_id=joined.trace_id, external_trace_id=joined.external_trace_id, parent=tool_node, trace_root=joined.root + ) diff --git a/api/domain/tracing/normalization.py b/api/domain/tracing/normalization.py index aba79110d456ff..e7a2c337fb7b53 100644 --- a/api/domain/tracing/normalization.py +++ b/api/domain/tracing/normalization.py @@ -278,10 +278,6 @@ def normalize_agent_span(data: Mapping[str, Any], *, capture_content: bool) -> S "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", } ) @@ -317,6 +313,7 @@ def remap_outer_trace( session_id=trace.session_id, user_id=trace.user_id, parent=None, + trace_root=None, spans=tuple(spans), ) @@ -363,6 +360,7 @@ def legacy_outer_trace(data: Mapping[str, Any]) -> Trace: session_id=data.get("session_id") or None, user_id=None, parent=None, + trace_root=None, spans=tuple(spans), ) diff --git a/api/domain/tracing/trace_run_agg.py b/api/domain/tracing/trace_run_agg.py index f836154cc726e4..0519f83f1acc72 100644 --- a/api/domain/tracing/trace_run_agg.py +++ b/api/domain/tracing/trace_run_agg.py @@ -422,6 +422,7 @@ def settle_preview(self, node: SourceNodeExecution, *, now: datetime) -> None: session_id=self.session_id, user_id=None, parent=None, + trace_root=None, spans=(span,), ) self._keep_terminal(trace, now=now) @@ -638,5 +639,6 @@ def exported_trace(self, batches: Sequence[TraceBatch]) -> Trace: session_id=self.session_id, user_id=outer.user_id, parent=None, + trace_root=None, spans=order_spans(outer.root.span_id, spans), ) diff --git a/api/extensions/application_services/tracing.py b/api/extensions/application_services/tracing.py index 26e8b3d61ecc9a..aaf0000d63e1df 100644 --- a/api/extensions/application_services/tracing.py +++ b/api/extensions/application_services/tracing.py @@ -46,6 +46,9 @@ def build_tracing_services(*, database_client: sessionmaker[Session]) -> Tracing builder = CanonicalTraceBuilder() configs = SQLAlchemyAppTracingConfigRepository(session_factory=database_client) enterprise_telemetry = is_enterprise_telemetry_enabled() + records = TraceRecordRepository( + database_client, files_url=dify_config.FILES_URL, enterprise_details=enterprise_telemetry + ) agent_settings = AgentTraceSettings( enabled=dify_config.OPS_AGENT_TRACE_ENABLED, max_spans=dify_config.OPS_AGENT_TRACE_MAX_SPANS, @@ -74,9 +77,8 @@ def build_tracing_services(*, database_client: sessionmaker[Session]) -> Tracing emitter=CeleryTraceEmitter(), policies=configs, enterprise_enabled=is_enterprise_telemetry_enabled ), delivery=TraceDeliveryService( - records=TraceRecordRepository( - database_client, files_url=dify_config.FILES_URL, enterprise_details=enterprise_telemetry - ), + records=records, + outer_runs=records, destinations=configs, providers=providers, codec=codec, diff --git a/api/extensions/tracing.py b/api/extensions/tracing.py index 93e60916e43063..a1512d499e0bb6 100644 --- a/api/extensions/tracing.py +++ b/api/extensions/tracing.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Literal, Protocol, cast from dify_trace_sdk.emitter import TraceScope -from dify_trace_sdk.events import TraceCorrelation +from dify_trace_sdk.events import TraceCorrelation, TraceParent from flask import Flask, current_app, has_app_context from core.helper.trace_id_helper import ParentTraceContext @@ -88,6 +88,17 @@ def _scopes() -> ExecutionTraceScopes | None: return cast(ExecutionTraceScopes | None, current_app.extensions.get(_SCOPES_KEY)) +def _trace_parent(context: ParentTraceContext) -> TraceParent: + outer = context.parent_trace_context + return TraceParent( + workflow_run_id=context.parent_workflow_run_id, + app_id=context.parent_app_id, + node_execution_id=context.parent_node_execution_id, + external_trace_id=context.parent_external_trace_id, + parent=_trace_parent(outer) if outer is not None 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") @@ -100,8 +111,7 @@ def correlation_from_extras(extras: Mapping[str, Any]) -> TraceCorrelation: 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, + parent=_trace_parent(parent) if isinstance(parent, ParentTraceContext) else None, ) 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 index b587bba7c8d837..b275b49f13a01c 100644 --- a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/provider.py +++ b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/provider.py @@ -118,7 +118,7 @@ def _observation_id(trace: Trace, span_id: str) -> str: def _events(self, trace: Trace, trace_key: str) -> list[dict[str, object]]: events: list[dict[str, object]] = [] - if trace.parent is None: + if trace.trace_root is None: root = trace.root events.append( self._event( @@ -140,7 +140,8 @@ def _events(self, trace: Trace, trace_key: str) -> list[dict[str, object]]: trace_level_span_id = root.span_id observations: Sequence[Span] = trace.spans[1:] else: - trace_level_span_id = trace.parent.span_id + # Spans under the trace root, which another unit sent, are top-level observations. + trace_level_span_id = trace.trace_root.span_id observations = trace.spans for span in observations: parent = ( 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 index 0c04bf5edb88cc..0b6227958915a6 100644 --- 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 @@ -11,7 +11,16 @@ 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 dify_trace_sdk.testing import ( + MESSAGE_ID, + WORKFLOW_RUN_ID, + RecordingEgress, + message_trace, + nested_workflow_fragment, + response, + tool_fragment, + workflow_trace, +) CONFIG = LangfuseConfig(public_key="pk-lf-1", secret_key="sk-lf-1", host="https://cloud.langfuse.com") @@ -100,6 +109,20 @@ def test_fragment_joins_the_message_trace_without_recreating_it() -> None: assert "parentObservationId" not in tool["body"] +def test_nested_workflow_hangs_from_the_tool_node_that_started_it() -> None: + provider, egress = _provider(response(207, {"successes": [{"id": "x", "status": 201}], "errors": []})) + fragment = nested_workflow_fragment() + assert fragment.parent is not None + + provider.send(CONFIG, fragment) + + run, node = _sent_events(egress) + assert [run["type"], node["type"]] == ["span-create", "span-create"] + assert run["body"]["traceId"] == node["body"]["traceId"] == WORKFLOW_RUN_ID + assert run["body"]["parentObservationId"] == str(span_uuid(fragment.trace_id, fragment.parent.span_id)) + assert node["body"]["parentObservationId"] == run["body"]["id"] + + 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"}) 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 index aa48e92110fd80..16b792e71a97c7 100644 --- a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/provider.py +++ b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/provider.py @@ -137,7 +137,7 @@ def _trace_run_id(trace: Trace) -> str: 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 + anchor = trace.trace_root.span_id if trace.trace_root is not None else trace.root.span_id return str(span_uuid(trace.trace_id, anchor)) @@ -251,12 +251,15 @@ def send(self, config: LangSmithConfig, trace: Trace) -> SendResult: 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: + if trace.parent is None or trace.trace_root 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) + # A deeper parent's dotted order names every run above it, so the unit hangs from the root. + root_order = _dotted_order(root_run_id, trace.trace_root.start_time, None) + for anchor in (trace.trace_root.span_id, trace.parent.span_id): + run_ids[anchor] = root_run_id + dotted_orders[anchor] = root_order client = self._client(config, api_url) accepted: list[str] = [] try: 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 index 10ed1402d3b1f7..c1fa8496cd7789 100644 --- 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 @@ -20,6 +20,7 @@ WORKFLOW_RUN_ID, RecordingEgress, message_trace, + nested_workflow_fragment, response, tool_fragment, workflow_trace, @@ -277,6 +278,25 @@ def test_fragment_attaches_under_the_message_run() -> None: ) +def test_nested_workflow_attaches_under_the_trace_root_run() -> None: + """LangSmith orders runs by their full ancestry, which a deeper parent's unit does not carry.""" + workflow_provider, workflow_egress = _provider(*_accepted(3)) + workflow_provider.send(VERIFIED, workflow_trace()) + root_order = _runs(workflow_egress)[0]["dotted_order"] + provider, egress = _provider(*_accepted(2)) + fragment = nested_workflow_fragment() + run_id = str(span_uuid(fragment.trace_id, fragment.root.span_id)) + + provider.send(VERIFIED, fragment) + + nested_run, node_run = _runs(egress) + assert nested_run["id"] == run_id + assert (nested_run["parent_run_id"], nested_run["trace_id"]) == (WORKFLOW_RUN_ID, WORKFLOW_RUN_ID) + assert nested_run["dotted_order"] == f"{root_order}.20261001T080003100000Z{run_id}" + assert node_run["parent_run_id"] == run_id + assert node_run["dotted_order"].startswith(f"{nested_run['dotted_order']}.") + + 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/") 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 index aee92de2426076..1ab12d34a84dce 100644 --- a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/_spans.py +++ b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/_spans.py @@ -55,6 +55,7 @@ def standalone(trace: Trace) -> Trace: update={ "trace_id": trace_id_for(f"{trace.trace_id}:{root.span_id}"), "parent": None, + "trace_root": None, "spans": (detached, *trace.spans[1:]), } ) 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 index 9624c813690de2..e1ea2132bbf565 100644 --- 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 @@ -28,6 +28,7 @@ RecordedRequest, RecordingEgress, message_trace, + nested_workflow_fragment, response, tool_fragment, workflow_trace, @@ -349,6 +350,22 @@ def test_fragment_is_sent_as_its_own_trace_linked_to_its_parent() -> None: assert result.trace_id == f"tr-{own_trace_id}" assert own_trace_id != fragment.trace_id + @staticmethod + def test_nested_workflow_is_its_own_trace_naming_the_tool_node_it_hangs_from() -> None: + provider, egress = _mlflow_provider(ACK) + fragment = nested_workflow_fragment() + assert fragment.parent is not None + own_trace_id = trace_id_for(f"{fragment.trace_id}:{fragment.root.span_id}") + + provider.send(MLFLOW_CONFIG, fragment) + + run, node = _exported(egress) + assert (run.trace_id, node.trace_id) == (f"tr-{own_trace_id}", f"tr-{own_trace_id}") + assert run.parent_id is None + assert run.get_attribute("dify.parent_trace_id") == fragment.trace_id + assert run.get_attribute("dify.parent_span_id") == fragment.parent.span_id + assert node.parent_id == run.span_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"}) 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 index a1f6de4a5dd229..4e1639629924f2 100644 --- a/api/providers/trace/trace-opik/src/dify_trace_opik/provider.py +++ b/api/providers/trace/trace-opik/src/dify_trace_opik/provider.py @@ -226,8 +226,8 @@ def send(self, config: OpikConfig, trace: Trace) -> SendResult: 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 + # A fragment names the root of its trace, so both units derive the same trace ID. + trace_start = trace.trace_root.start_time if trace.trace_root 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) @@ -278,6 +278,9 @@ def _check(self, method: str, url: str, headers: Mapping[str, str], content: byt 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} + if trace.parent is not None and trace.parent != trace.trace_root: + # A parent below the trace root is a span another unit sent; its ID is derivable. + span_ids[trace.parent.span_id] = _opik_id(trace.trace_id, trace.parent.span_id, trace.parent.start_time) 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: @@ -312,7 +315,7 @@ def _span(trace: Trace, span: Span, trace_id: str, span_ids: Mapping[str, str], "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. + # The unit's root is top-level, and so is a fragment's root that hangs from the trace root. "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"), 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 index 00a9890e8f6022..a968a87f0de71c 100644 --- 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 @@ -24,6 +24,7 @@ RecordedRequest, RecordingEgress, message_trace, + nested_workflow_fragment, response, tool_fragment, workflow_trace, @@ -456,6 +457,7 @@ def test_large_units_send_spans_in_batches() -> None: session_id=None, user_id=None, parent=None, + trace_root=None, spans=(run, *nodes), ) provider, egress = _provider(CREATED, NO_CONTENT, NO_CONTENT) @@ -514,6 +516,20 @@ def test_span_failures_after_the_trace_was_created_are_outcome_unknown(outcome: assert len(egress.requests) == 2 +def test_nested_workflow_hangs_from_the_tool_node_span_of_its_outer_trace() -> None: + provider, egress = _provider(NO_CONTENT) + + assert provider.send(CONFIG, nested_workflow_fragment()).trace_id == WORKFLOW_TRACE_ID + + (create_spans,) = egress.requests + assert create_spans.url == f"{API_URL}v1/private/spans/batch" + body = create_spans.json() + assert isinstance(body, dict) + run, node = body["spans"] + assert (run["trace_id"], run["parent_span_id"]) == (WORKFLOW_TRACE_ID, WORKFLOW_TOOL_SPAN_ID) + assert (node["trace_id"], node["parent_span_id"]) == (WORKFLOW_TRACE_ID, run["id"]) + + @pytest.mark.parametrize( ("outcome", "expected"), [(response(503, {}), ExportRetryableError), (response(400, {}), ExportRejectedError)], 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 index 3f1b8dd3ac3714..4a55503b75f94b 100644 --- a/api/providers/trace/trace-sdk/src/dify_trace_sdk/events.py +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/events.py @@ -20,13 +20,29 @@ def json_safe(value: object) -> JsonValue: return json.loads(json.dumps(value, ensure_ascii=False, default=str), parse_constant=str) +class TraceParent(BaseModel): + """The workflow run, and its tool node execution, that started a nested workflow run.""" + + workflow_run_id: str + app_id: str | None + """``None`` when the context that started the nested run did not record it.""" + node_execution_id: str | None + """``None`` when the starting node execution is unknown.""" + external_trace_id: str | None + """The trace ID that run's caller supplied; only an outermost run has one.""" + parent: "TraceParent | None" + """The run that started that run in turn; ``None`` when that run is not nested.""" + + model_config = _EVENT_CONFIG + + 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 + parent: TraceParent | None + """The run that started this nested workflow run; ``None`` outside nested workflows.""" model_config = _EVENT_CONFIG 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 index 3203233f084595..adfe42c630725a 100644 --- a/api/providers/trace/trace-sdk/src/dify_trace_sdk/model.py +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/model.py @@ -153,8 +153,10 @@ 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. + another unit: a tool call under its message, or a nested workflow run under + the tool node that started it. The first span of this unit points at it. + ``trace_root`` is the root span of that trace, the same span as ``parent`` + when the unit hangs directly from it. """ trace_id: str @@ -163,6 +165,8 @@ class Trace(BaseModel): session_id: str | None user_id: str | None parent: SpanRef | None + trace_root: SpanRef | None + """Set exactly when ``parent`` is set.""" spans: tuple[Span, ...] = Field(min_length=1) model_config = ConfigDict(extra="forbid", frozen=True) @@ -187,6 +191,10 @@ def _parent_first(self) -> Self: seen.add(span.span_id) if expected_root_parent in seen: raise ValueError("the external parent cannot be a span of this unit") + if (self.parent is None) != (self.trace_root is None): + raise ValueError("parent and trace_root must be set together") + if self.trace_root is not None and self.trace_root.span_id in seen: + raise ValueError("the trace root cannot be a span of this unit") return self @property 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 index 846d38716062a4..60fde5fb60f468 100644 --- a/api/providers/trace/trace-sdk/src/dify_trace_sdk/testing.py +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/testing.py @@ -15,6 +15,7 @@ MESSAGE_ID = "7b3c6c0e-38a1-4b52-9a7e-3ad8f3a1c001" WORKFLOW_RUN_ID = "7b3c6c0e-38a1-4b52-9a7e-3ad8f3a1c002" +NESTED_WORKFLOW_RUN_ID = "7b3c6c0e-38a1-4b52-9a7e-3ad8f3a1c003" STARTED_AT = datetime(2026, 10, 1, 8, 0, 0, tzinfo=UTC) @@ -89,6 +90,7 @@ def message_trace() -> Trace: session_id="conversation-1", user_id="end-user-1", parent=None, + trace_root=None, spans=( Span( span_id=message_span_id, @@ -133,6 +135,7 @@ def workflow_trace() -> Trace: session_id=WORKFLOW_RUN_ID, user_id="account-1", parent=None, + trace_root=None, spans=( Span( span_id=run_span_id, @@ -187,6 +190,7 @@ def tool_fragment() -> Trace: session_id="conversation-1", user_id="end-user-1", parent=SpanRef(span_id=span_id_for(MESSAGE_ID), start_time=STARTED_AT), + trace_root=SpanRef(span_id=span_id_for(MESSAGE_ID), start_time=STARTED_AT), spans=( Span( span_id=span_id_for("tool-event-1"), @@ -202,3 +206,47 @@ def tool_fragment() -> Trace: ), ), ) + + +def nested_workflow_fragment() -> Trace: + """A nested workflow run sent separately and attached under the tool node of ``workflow_trace``.""" + tool_span_id = span_id_for("node-execution-tool") + run_span_id = span_id_for(NESTED_WORKFLOW_RUN_ID) + return Trace( + trace_id=trace_id_for(WORKFLOW_RUN_ID), + external_trace_id=None, + session_id=WORKFLOW_RUN_ID, + user_id="account-1", + parent=SpanRef(span_id=tool_span_id, start_time=STARTED_AT + timedelta(seconds=3)), + trace_root=SpanRef(span_id=span_id_for(WORKFLOW_RUN_ID), start_time=STARTED_AT), + spans=( + Span( + span_id=run_span_id, + parent_span_id=tool_span_id, + name="workflow", + kind=SpanKind.CHAIN, + start_time=STARTED_AT + timedelta(seconds=3, milliseconds=100), + end_time=STARTED_AT + timedelta(seconds=4), + status=SpanStatus.OK, + inputs={"query": "tracing"}, + outputs={"answer": "Tracing records execution."}, + attributes={ + "workflow_run_id": NESTED_WORKFLOW_RUN_ID, + "parent_workflow_run_id": WORKFLOW_RUN_ID, + "parent_node_execution_id": "node-execution-tool", + }, + ), + Span( + span_id=span_id_for("nested-node-execution-code"), + parent_span_id=run_span_id, + name="code_Lookup", + kind=SpanKind.CHAIN, + start_time=STARTED_AT + timedelta(seconds=3, milliseconds=200), + end_time=STARTED_AT + timedelta(seconds=3, milliseconds=900), + status=SpanStatus.OK, + inputs={"query": "tracing"}, + outputs={"answer": "Tracing records execution."}, + attributes={"node_execution_id": "nested-node-execution-code", "node_type": "code"}, + ), + ), + ) 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 index 67b5b7b56afc2d..42d5428173b154 100644 --- 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 @@ -4,9 +4,7 @@ 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 -) +CORRELATION = TraceCorrelation(external_trace_id="caller-1", session_id=None, parent=None) class _CollectingEmitter: 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 index a66539abd6a8d6..ad7f57dfe31253 100644 --- 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 @@ -59,6 +59,7 @@ def test_trace_requires_parent_first_order() -> None: session_id=None, user_id=None, parent=None, + trace_root=None, spans=(_span(root, None), _span(span_id_for("orphan"), child)), ) with pytest.raises(ValidationError, match="root of this unit"): @@ -68,6 +69,7 @@ def test_trace_requires_parent_first_order() -> None: session_id=None, user_id=None, parent=None, + trace_root=None, spans=(_span(child, root),), ) @@ -81,10 +83,36 @@ def test_fragment_root_must_point_at_its_external_parent() -> None: session_id=None, user_id=None, parent=parent, + trace_root=parent, spans=(_span(span_id_for("tool"), None),), ) +def test_fragment_names_its_parent_and_the_root_of_its_trace_together() -> None: + trace_root = SpanRef(span_id=span_id_for("workflow"), start_time=STARTED_AT) + tool_node = SpanRef(span_id=span_id_for("tool-node"), start_time=STARTED_AT) + nested = _span(span_id_for("nested-workflow"), tool_node.span_id) + + def fragment(parent: SpanRef | None, root: SpanRef | None, spans: tuple[Span, ...]) -> Trace: + return Trace( + trace_id=trace_id_for("t"), + external_trace_id=None, + session_id=None, + user_id=None, + parent=parent, + trace_root=root, + spans=spans, + ) + + assert fragment(tool_node, trace_root, (nested,)).trace_root == trace_root + with pytest.raises(ValidationError, match="set together"): + fragment(tool_node, None, (nested,)) + with pytest.raises(ValidationError, match="set together"): + fragment(None, trace_root, (_span(span_id_for("root"), None),)) + with pytest.raises(ValidationError, match="trace root cannot be a span of this unit"): + fragment(tool_node, SpanRef(span_id=nested.span_id, start_time=STARTED_AT), (nested,)) + + def test_cost_must_be_finite() -> None: with pytest.raises(ValidationError): Cost(amount=float("inf"), currency="USD") 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 index c72e06a793e7d8..f54c457d55aa85 100644 --- 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 @@ -10,7 +10,7 @@ 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 dify_trace_sdk.testing import RecordingEgress, message_trace, nested_workflow_fragment, 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 @@ -26,6 +26,18 @@ def test_otlp_encoding_keeps_ids_and_external_parent() -> None: assert [attribute.key for attribute in span.attributes] == ["a"] +def test_otlp_encoding_hangs_a_nested_unit_from_the_span_that_started_it() -> None: + fragment = nested_workflow_fragment() + assert fragment.parent is not None + + request = encode_trace(fragment, resource={"service.name": "dify"}, attributes=lambda span: {}) + + run, node = request.resource_spans[0].scope_spans[0].spans + assert run.trace_id.hex() == node.trace_id.hex() == fragment.trace_id + assert run.parent_span_id.hex() == fragment.parent.span_id + assert node.parent_span_id.hex() == fragment.root.span_id + + def test_otlp_ack_partial_success_is_outcome_unknown() -> None: partial = ExportTraceServiceResponse() partial.partial_success.rejected_spans = 1 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 index 7a97ac9c91c363..d7b86dfc91e534 100644 --- 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 @@ -1,4 +1,4 @@ -from dify_trace_sdk.testing import message_trace, tool_fragment, workflow_trace +from dify_trace_sdk.testing import message_trace, nested_workflow_fragment, tool_fragment, workflow_trace def test_sample_traces_are_valid_parent_first_trees() -> None: @@ -7,3 +7,10 @@ def test_sample_traces_are_valid_parent_first_trees() -> None: fragment = tool_fragment() assert fragment.parent is not None assert fragment.root.parent_span_id == fragment.parent.span_id + nested = nested_workflow_fragment() + assert nested.parent is not None + assert nested.trace_root is not None + assert nested.root.parent_span_id == nested.parent.span_id + assert nested.trace_id == workflow_trace().trace_id + assert nested.trace_root.span_id == workflow_trace().root.span_id + assert nested.parent.span_id in {span.span_id for span in workflow_trace().spans} 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 index 9554d3bf201a20..d962fd1de58c18 100644 --- a/api/providers/trace/trace-weave/src/dify_trace_weave/provider.py +++ b/api/providers/trace/trace-weave/src/dify_trace_weave/provider.py @@ -281,7 +281,7 @@ def send(self, config: WeaveConfig, trace: Trace) -> SendResult: 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 + root_span_id = trace.trace_root.span_id if trace.trace_root 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)}", 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 index 6253d07e071a19..a065a4e592fa2c 100644 --- 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 @@ -22,6 +22,7 @@ RecordedRequest, RecordingEgress, message_trace, + nested_workflow_fragment, response, tool_fragment, workflow_trace, @@ -477,6 +478,22 @@ def test_fragment_joins_the_message_call_without_naming_a_turn() -> None: ) +def test_nested_workflow_hangs_from_the_tool_node_call_that_started_it() -> None: + fragment = nested_workflow_fragment() + assert fragment.parent is not None + assert fragment.trace_root is not None + provider, egress = _provider(*_acknowledgments(fragment)) + + result = provider.send(CONFIG, fragment) + + run_start, node_start = (_body(request)["start"] for request in egress.requests[:2]) + assert (run_start["trace_id"], node_start["trace_id"]) == (WORKFLOW_RUN_ID, WORKFLOW_RUN_ID) + assert run_start["parent_id"] == _call_id(fragment, fragment.parent.span_id) + assert node_start["parent_id"] == run_start["id"] + assert run_start["turn_id"] is None + assert result.trace_url == f"{CALLS_URL}/{_call_id(fragment, fragment.trace_root.span_id)}" + + 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)) diff --git a/api/repositories/trace_record_repository.py b/api/repositories/trace_record_repository.py index 0b6bd57a1f3522..9a19a994995e05 100644 --- a/api/repositories/trace_record_repository.py +++ b/api/repositories/trace_record_repository.py @@ -31,13 +31,13 @@ from graphon.enums import WorkflowExecutionStatus from models import Account, App, Tenant from models.dataset import Dataset -from models.model import Conversation, EndUser, Message, MessageFile +from models.model import AppMode, 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, WorkflowNodeExecutionTriggeredFrom, WorkflowRun from repositories.factory import DifyAPIRepositoryFactory from repositories.tracing_store_errors import tracing_store_errors -from services.tracing.delivery import TraceRecordSource +from services.tracing.delivery import OuterRunQuery, OuterRunRoot, TraceRecordSource from services.tracing.errors import SourceNotFoundError, SourceNotReadyError from services.tracing.records import ( BaseTraceInfo, @@ -123,7 +123,7 @@ def _streaming_metrics(message: Message) -> dict[str, Any]: } -class TraceRecordRepository(TraceRecordSource): +class TraceRecordRepository(TraceRecordSource, OuterRunQuery): def __init__(self, session_factory: sessionmaker[Session], *, files_url: str, enterprise_details: bool) -> None: self._sessions = session_factory self._files_url = files_url @@ -154,6 +154,38 @@ def assemble(self, event: TraceEvent) -> BaseTraceInfo: # Shared lookups + @override + def root(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> OuterRunRoot | None: + runs = DifyAPIRepositoryFactory.create_api_workflow_run_repository(self._sessions) + with tracing_store_errors(): + run = runs.get_workflow_run_by_id(tenant_id=tenant_id, app_id=app_id, run_id=workflow_run_id) + if run is None: + return None + with tracing_store_errors(), self._sessions() as session: + mode = session.scalar(select(App.mode).where(App.id == app_id, App.tenant_id == tenant_id)) + if mode is not AppMode.ADVANCED_CHAT: + return OuterRunRoot(root_id=run.id, started_at=run.created_at) + message = session.execute( + select(Message.id, Message.created_at).where( + Message.app_id == app_id, Message.workflow_run_id == run.id + ) + ).first() + if message is None: + return None + message_id, created_at = message + return OuterRunRoot(root_id=message_id, started_at=created_at) + + @override + def node_started_at( + self, *, tenant_id: str, app_id: str, workflow_run_id: str, node_execution_id: str + ) -> datetime | None: + executions = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(self._sessions) + with tracing_store_errors(): + execution = executions.get_execution_by_id(execution_id=node_execution_id, tenant_id=tenant_id) + if execution is None or execution.app_id != app_id or execution.workflow_run_id != workflow_run_id: + return None + return execution.created_at + 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: @@ -330,10 +362,10 @@ def _workflow(self, event: TraceEvent, payload: WorkflowRunCompleted) -> Workflo "workspace_name": workspace_name, } correlation = event.correlation - if correlation.parent_workflow_run_id: + if correlation.parent is not None: metadata["parent_trace_context"] = { - "parent_workflow_run_id": correlation.parent_workflow_run_id, - "parent_node_execution_id": correlation.parent_node_execution_id, + "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, @@ -360,6 +392,7 @@ def _workflow(self, event: TraceEvent, payload: WorkflowRunCompleted) -> Workflo start_time=run.created_at, end_time=run.finished_at, invoked_by=_user_of(metadata), + parent=correlation.parent, ) def _node_executions(self, run: WorkflowRun) -> Sequence[WorkflowNodeExecution]: diff --git a/api/services/tracing/canonical.py b/api/services/tracing/canonical.py index 505fb01a7a645b..83806c5ffe82be 100644 --- a/api/services/tracing/canonical.py +++ b/api/services/tracing/canonical.py @@ -43,8 +43,9 @@ _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 root_trace_id(external_trace_id: str | None, root_id: str) -> str: + """The trace ID of a trace rooted at the message or workflow run ``root_id``.""" + return trace_id_for(external_trace_id or root_id) def _status(error: str | None, status: object) -> SpanStatus: @@ -179,11 +180,12 @@ def _message(self, record: MessageTraceInfo) -> Trace: status = _status(error, None) attributes = _attributes({**record.metadata, "trace_entity_type": "message"}) return Trace( - trace_id=message_trace_id(record.trace_id, message.id), + trace_id=root_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, + trace_root=None, spans=( Span( span_id=root_span_id, @@ -323,11 +325,12 @@ def _workflow(self, record: WorkflowTraceInfo) -> Trace: } ) return Trace( - trace_id=trace_id_for(record.trace_id or root_logical_id), + trace_id=root_trace_id(record.trace_id, 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, + trace_root=None, spans=tuple(spans), ) @@ -416,7 +419,7 @@ def _auxiliary( parent = None parent_span_id = None else: - trace_id = message_trace_id(record.trace_id, message.id) + trace_id = root_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( @@ -425,6 +428,7 @@ def _auxiliary( session_id=_session_id(record, message.conversation_id if message is not None else None), user_id=_user_id(record), parent=parent, + trace_root=parent, spans=( Span( span_id=span_id, diff --git a/api/services/tracing/delivery.py b/api/services/tracing/delivery.py index 9e91e3e0a80d22..d24bb75e8827dc 100644 --- a/api/services/tracing/delivery.py +++ b/api/services/tracing/delivery.py @@ -2,31 +2,36 @@ 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. +tracing provider receives the rest. A nested workflow run joins the trace of +the run that started it when both go to the same destination. 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 collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass +from datetime import datetime 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.events import TraceEvent, TraceParent +from dify_trace_sdk.ids import span_id_for +from dify_trace_sdk.model import SpanRef, 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 domain.tracing.nested_workflow_trace import OuterTrace, nested_trace_placement +from services.tracing.canonical import CanonicalTraceBuilder, root_trace_id from services.tracing.errors import ( ProviderConfigUnreadableError, TraceCollectorUnavailableError, TraceProviderUnavailableError, + TracingStoreUnavailableError, UnknownTraceProviderError, ) from services.tracing.providers import ProviderConfigCodec, TraceProviderCatalog @@ -67,13 +72,46 @@ def active_destination(self, *, tenant_id: str, app_id: str) -> StoredTraceDesti ... +@dataclass(frozen=True, slots=True) +class OuterRunRoot: + """The span a workflow run's trace starts with when that run is the root of its trace.""" + + root_id: str + """The chatflow message the run answers, otherwise the run itself.""" + started_at: datetime + + +class OuterRunQuery(Protocol): + """The runs that started nested workflow runs, read through their tenant and app.""" + + def root(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> OuterRunRoot | None: + """``None`` until the run, and the message a chatflow run answers, are persisted.""" + ... + + def node_started_at( + self, *, tenant_id: str, app_id: str, workflow_run_id: str, node_execution_id: str + ) -> datetime | None: + """``None`` until the node execution is persisted.""" + ... + + +@dataclass(frozen=True, slots=True) +class CollectedTrace: + """How the Agent trace of a workflow run is exported.""" + + trace_id: str | None + """``None`` until collection starts and fixes it.""" + destination: StoredTraceDestination | None + """The destination frozen when collection started; ``None`` when the run exports nowhere.""" + + 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.""" + def collected_trace(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> CollectedTrace | None: + """How the run's Agent trace is exported; ``None`` when its spans are not collected.""" ... @@ -103,6 +141,10 @@ class Destination: provider: TraceProvider[Any] config: ProviderConfig + def is_same(self, other: "Destination") -> bool: + """Whether both send to the same provider with the same configuration.""" + return self.provider.id == other.provider.id and self.config == other.config + def retry_delay_seconds(attempt: int) -> int: return min(600, 10 * 2 ** (attempt - 1)) @@ -113,6 +155,7 @@ def __init__( self, *, records: TraceRecordSource, + outer_runs: OuterRunQuery, destinations: TraceDestinationQuery, providers: TraceProviderCatalog, codec: ProviderConfigCodec, @@ -125,6 +168,7 @@ def __init__( ) -> None: """``enterprise_sink`` and ``collector`` are ``None`` when the deployment does not run that sink.""" self._records = records + self._outer_runs = outer_runs self._destinations = destinations self._providers = providers self._codec = codec @@ -154,8 +198,6 @@ def deliver(self, event: TraceEvent, *, enterprise_delivered: bool) -> None: 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: @@ -178,6 +220,8 @@ def deliver(self, event: TraceEvent, *, enterprise_delivered: bool) -> None: if trace is None: _event_counter.add(1, {"kind": kind, "outcome": "not_exported"}) return + if isinstance(record, WorkflowTraceInfo) and record.parent is not None: + trace = self._nested(event.tenant_id, record.parent, destination, trace) _event_counter.add(1, {"kind": kind, "outcome": "exported"}) self.export( ExportJob( @@ -232,8 +276,71 @@ def export(self, job: ExportJob, destination: Destination) -> None: 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 + return self._decoded(tenant_id, app_id, stored) if stored is not None else None + + def _nested(self, tenant_id: str, parent: TraceParent, destination: Destination, trace: Trace) -> Trace: + """Hang a nested run's trace from the tool node that started it when its outer trace goes to ``destination``. + + A nested run whose outer records cannot be read yet is sent on its own. + """ + if parent.app_id is None or parent.node_execution_id is None: + return trace + try: + started_at = self._outer_runs.node_started_at( + tenant_id=tenant_id, + app_id=parent.app_id, + workflow_run_id=parent.workflow_run_id, + node_execution_id=parent.node_execution_id, + ) + if started_at is None: + return trace + placement = nested_trace_placement( + self._outer_traces(tenant_id, parent, destination), + tool_node=SpanRef(span_id=span_id_for(parent.node_execution_id), start_time=started_at), + ) + except TracingStoreUnavailableError: + logger.warning("Outer trace unreadable; nested trace sent on its own: app_id=%s", parent.app_id) + return trace + return placement.place(trace) if placement is not None else trace + + def _outer_traces(self, tenant_id: str, parent: TraceParent, destination: Destination) -> Iterator[OuterTrace]: + """The traces of the runs that started a nested run, innermost first, while they are known.""" + link: TraceParent | None = parent + while link is not None and link.app_id is not None: + root = self._outer_runs.root(tenant_id=tenant_id, app_id=link.app_id, workflow_run_id=link.workflow_run_id) + if root is None: + return + collected = ( + self._collector.collected_trace( + tenant_id=tenant_id, app_id=link.app_id, workflow_run_id=link.workflow_run_id + ) + if self._collector is not None + else None + ) + if collected is None: + trace_id = root_trace_id(link.external_trace_id, root.root_id) + external_trace_id = link.external_trace_id + outer_destination = self.destination(tenant_id, link.app_id) + elif collected.trace_id is None: + return + else: + # An Agent trace keeps a caller's trace ID only as an attribute. + trace_id, external_trace_id = collected.trace_id, None + outer_destination = ( + self._decoded(tenant_id, link.app_id, collected.destination) + if collected.destination is not None + else None + ) + yield OuterTrace( + trace_id=trace_id, + external_trace_id=external_trace_id, + root=SpanRef(span_id=span_id_for(root.root_id), start_time=root.started_at), + same_destination=outer_destination is not None and outer_destination.is_same(destination), + collected=collected is not None, + ) + link = link.parent + + def _decoded(self, tenant_id: str, app_id: str, stored: StoredTraceDestination) -> Destination | None: try: provider = self._providers.resolve(stored.provider_id) config = self._codec.decode(tenant_id, provider, stored.config) diff --git a/api/services/tracing/records.py b/api/services/tracing/records.py index 0e1a911a916304..448ca0e882e1d6 100644 --- a/api/services/tracing/records.py +++ b/api/services/tracing/records.py @@ -11,6 +11,7 @@ from enum import StrEnum from typing import Any +from dify_trace_sdk.events import TraceParent from pydantic import BaseModel, ConfigDict, StrictStr, ValidationError from graphon.entities import WorkflowNodeExecution @@ -114,6 +115,8 @@ class WorkflowTraceInfo(BaseTraceInfo): file_list: list[str] invoked_by: str | None query: str + parent: TraceParent | None + """The run that started this nested workflow run; ``None`` outside nested workflows.""" class MessageTraceInfo(BaseTraceInfo): diff --git a/api/services/tracing/recovery.py b/api/services/tracing/recovery.py index 18e69ff0c16234..a3fa548b0f7cd6 100644 --- a/api/services/tracing/recovery.py +++ b/api/services/tracing/recovery.py @@ -39,8 +39,7 @@ def _event(run: TraceRunAgg, payload: TracePayload, user_id: str | None) -> Trac correlation=TraceCorrelation( external_trace_id=None, session_id=run.session_id, - parent_workflow_run_id=None, - parent_node_execution_id=None, + parent=None, ), payload=payload, ) diff --git a/api/services/tracing/scopes.py b/api/services/tracing/scopes.py index 86f3a2edbe9f1f..1652eafa32955b 100644 --- a/api/services/tracing/scopes.py +++ b/api/services/tracing/scopes.py @@ -16,9 +16,7 @@ logger = logging.getLogger(__name__) -UNCORRELATED = TraceCorrelation( - external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None -) +UNCORRELATED = TraceCorrelation(external_trace_id=None, session_id=None, parent=None) """The correlation of an operation that carries no caller-supplied trace context.""" diff --git a/api/services/tracing/service.py b/api/services/tracing/service.py index 39c1a8fcfd827a..25a2d267262cf4 100644 --- a/api/services/tracing/service.py +++ b/api/services/tracing/service.py @@ -40,7 +40,7 @@ workflow_has_agent_v2, ) from services.tracing.collection_ports import ScheduleTraceExport, TraceRunQuery, TraceRunRepository, TraceSourceQuery -from services.tracing.delivery import TraceCollector, TraceRecordSource +from services.tracing.delivery import CollectedTrace, StoredTraceDestination, TraceCollector, TraceRecordSource from services.tracing.errors import ( ProviderConfigUnreadableError, TraceCollectionConflictError, @@ -337,8 +337,7 @@ def capture_agent_app_terminal( 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, + parent=None, ), payload=MessageCompleted(message_id=message_id, conversation_id=source.conversation_id), ) @@ -432,24 +431,20 @@ def _managed_owner( return False @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 info - parent = self._run_query.workflow_run_trace(tenant_id=tenant_id, workflow_run_id=parent_run_id) - if parent is None: - 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 - if parent_node_id: - 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 - return info.model_copy(update={"metadata": metadata}) + def collected_trace(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> CollectedTrace | None: + run = self._run_query.workflow_run_trace(tenant_id=tenant_id, workflow_run_id=workflow_run_id) + if run is not None and run.app_id != app_id: + # A collected run of another app is not the run that started the nested run. + return CollectedTrace(trace_id=None, destination=None) + if run is not None: + destination = ( + StoredTraceDestination(provider_id=run.provider, config=run.provider_config) if run.provider else None + ) + return CollectedTrace(trace_id=run.trace_id, destination=destination) + app = self._sources.app(tenant_id=tenant_id, app_id=app_id) + if app is None or not self._managed_owner(app, None, workflow_run_id): + return None + return CollectedTrace(trace_id=None, destination=None) @override def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: 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 72afdc14757a3d..83eb9a52c21fff 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 @@ -66,9 +66,7 @@ def _trace_scope() -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) 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 49e56e36fb9ed3..e56c452c67aefc 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 @@ -100,9 +100,7 @@ def _trace_scope() -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) 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 9f2f0bafc60747..6a2cd410032cf8 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 @@ -93,9 +93,7 @@ def _trace_scope() -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) @@ -1666,9 +1664,7 @@ def raise_conversation_not_exists(**_kwargs): 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) scope_requests: list[dict[str, object]] = [] 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 db51ec30017348..bf7ba67f3f5c57 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 @@ -95,8 +95,7 @@ def _trace_scope(recorder: _TraceRecorder) -> TraceScope: correlation=TraceCorrelation( external_trace_id=None, session_id="trace-session-1", - parent_workflow_run_id=None, - parent_node_execution_id=None, + parent=None, ), ) @@ -437,9 +436,7 @@ def _timeline_trace_scope(recorder: _TimelineTraceRecorder) -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) 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 5c426513035a6b..29658525afe6e5 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 @@ -85,9 +85,7 @@ def _trace_scope() -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) @@ -613,9 +611,7 @@ def _make_entity(query: str) -> AgentAppGenerateEntity: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ), agent_id="agent-1", agent_config_snapshot_id="snapshot-1", 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 2bb1fdc1caaf7a..bbee5286394885 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 @@ -60,9 +60,7 @@ def _trace_scope() -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id="session-1", parent=None), ) 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 bfaf28a79b63a9..633efafaaae3b7 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 @@ -40,9 +40,7 @@ def _trace_scope() -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id="session-1", parent=None), ) 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 e648c14ab2aece..b09292b25b09f1 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 @@ -36,9 +36,7 @@ def _trace_scope() -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id="session-1", parent=None), ) 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 443caab0770093..0a4021ee5b6156 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 @@ -889,9 +889,7 @@ def check(self: object, **kwargs: object) -> tuple[bool, dict[str, object], str] 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) app_generate_entity = AppGenerateEntity( task_id="task", diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py index 4ffc0bde565467..60a9a4f94755ae 100644 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py @@ -48,6 +48,7 @@ ) from core.credit_usage import CreditUsageAppType from core.db.session_factory import session_factory +from core.helper.trace_id_helper import ParentTraceContext from core.repositories.human_input_repository import HumanInputFormRepositoryImpl, HumanInputFormSubmissionRepository from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository @@ -344,10 +345,19 @@ def fake_from_graph_init_context(**kwargs): invoke_from=InvokeFrom.DEBUGGER, root_node_id="root", trace_session_id="session-1", + external_trace_id="caller-trace", + parent_trace_context=ParentTraceContext( + parent_workflow_run_id="outer-run", parent_node_execution_id="outer-node-execution" + ), call_depth=3, ) - assert captured["run_context"][DIFY_RUN_CONTEXT_KEY].trace_session_id == "session-1" + dify_context = captured["run_context"][DIFY_RUN_CONTEXT_KEY] + assert dify_context.trace_session_id == "session-1" + assert dify_context.external_trace_id == "caller-trace" + assert dify_context.parent_trace_context == ParentTraceContext( + parent_workflow_run_id="outer-run", parent_node_execution_id="outer-node-execution" + ) assert captured["call_depth"] == 3 assert captured["runtime_state"] is runtime_state @@ -1180,6 +1190,8 @@ def test_run_uses_single_node_execution_branch( user_id="user", app_type=CreditUsageAppType.WORKFLOW, trace_session_id="session-1", + external_trace_id=None, + parent_trace_context=None, ) init_graph.assert_not_called() 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 1b8896b29d27fe..2bea85a00aad14 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 @@ -31,6 +31,7 @@ _WorkflowGenerateEntityWrapper, ) from core.app.workflow.file_runtime import DifyWorkflowFileRuntime +from core.helper.trace_id_helper import ParentTraceContext from core.repositories import SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from core.tools.tool_file_manager import ToolFileManager @@ -83,9 +84,7 @@ def _trace_scope(recorder: _TraceRecorder | _PersistedRunRecorder, *, app: App, 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) @@ -399,10 +398,9 @@ def fake_generate(**kwargs): assert isinstance(application_generate_entity, WorkflowAppGenerateEntity) extras = application_generate_entity.extras assert extras["external_trace_id"] == "trace-1" - assert extras["parent_trace_context"].model_dump() == { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - } + assert extras["parent_trace_context"] == ParentTraceContext( + parent_workflow_run_id="outer-workflow-run-1", 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 == [ @@ -718,9 +716,7 @@ def _app_generator_extra_trace_scope() -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) 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 8603ee44e8ad02..39af65e7bd9cba 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 @@ -40,9 +40,7 @@ def _trace_scope() -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) 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 3f5c8bc1295e25..94fc4183a358d1 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 @@ -455,9 +455,7 @@ def test_message_end_emits_message_trace_after_commit( 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) queue_manager.messages = [queue_message(QueueMessageEndEvent(llm_result=None))] saved_sessions: list[Session] = [] @@ -562,9 +560,7 @@ def _trace_scope(recorder: _TimelineTraceRecorder) -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) 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 13c5e78e74b727..a0893ac1c56d68 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 @@ -215,8 +215,7 @@ def test_generate_conversation_name_success( correlation=TraceCorrelation( external_trace_id="caller-trace", session_id=None, - parent_workflow_run_id=None, - parent_node_execution_id=None, + parent=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 ad51d7bbc1489d..cf7f51c0b7b2c5 100644 --- a/api/tests/unit_tests/core/moderation/test_input_moderation.py +++ b/api/tests/unit_tests/core/moderation/test_input_moderation.py @@ -26,9 +26,7 @@ def _trace_scope(recorder: _Recorder) -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) 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 2dfb989a6367a9..d7655a6e1c9bd7 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 @@ -98,9 +98,7 @@ def _trace_scope(recorder: _Recorder) -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) diff --git a/api/tests/unit_tests/core/telemetry/test___init__.py b/api/tests/unit_tests/core/telemetry/test___init__.py index 2692242d087127..af6799d0a5093f 100644 --- a/api/tests/unit_tests/core/telemetry/test___init__.py +++ b/api/tests/unit_tests/core/telemetry/test___init__.py @@ -33,9 +33,7 @@ def create_scope( 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) def emit(self, event: TraceEvent) -> None: diff --git a/api/tests/unit_tests/core/telemetry/test_gateway.py b/api/tests/unit_tests/core/telemetry/test_gateway.py index 6df19f851f74a3..4b39504ff353f9 100644 --- a/api/tests/unit_tests/core/telemetry/test_gateway.py +++ b/api/tests/unit_tests/core/telemetry/test_gateway.py @@ -120,9 +120,7 @@ def create_scope( 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) def emit(self, event: TraceEvent) -> None: 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 54abd16923187e..3ead7cdad779c3 100644 --- a/api/tests/unit_tests/core/tools/test_tool_engine.py +++ b/api/tests/unit_tests/core/tools/test_tool_engine.py @@ -57,9 +57,7 @@ def _trace_scope(recorder: _Recorder) -> TraceScope: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) diff --git a/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py b/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py index d066c8fe771fb5..109639d4776e81 100644 --- a/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py +++ b/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py @@ -13,6 +13,7 @@ from sqlalchemy.orm import Session, sessionmaker from core.app.entities.app_invoke_entities import InvokeFrom +from core.helper.trace_id_helper import ParentTraceContext from core.tools.__base.tool_runtime import ToolRuntime from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ( @@ -339,8 +340,15 @@ def test_workflow_tool_passes_parent_trace_context_from_runtime( """Ensure nested workflow runtime metadata is forwarded as parent trace context.""" tool = _build_tool() tool.set_parent_trace_context( - parent_workflow_run_id="outer-workflow-run-1", - parent_node_execution_id="outer-node-execution-1", + ParentTraceContext( + parent_workflow_run_id="outer-workflow-run-1", + parent_node_execution_id="outer-node-execution-1", + parent_trace_context=ParentTraceContext( + parent_workflow_run_id="outermost-workflow-run", + parent_node_execution_id="outermost-node-execution", + parent_external_trace_id="caller-trace", + ), + ) ) monkeypatch.setattr(tool, "_get_app", lambda *args, **kwargs: None) @@ -358,6 +366,15 @@ def test_workflow_tool_passes_parent_trace_context_from_runtime( assert call_kwargs["args"]["parent_trace_context"].model_dump() == { "parent_workflow_run_id": "outer-workflow-run-1", "parent_node_execution_id": "outer-node-execution-1", + "parent_app_id": None, + "parent_external_trace_id": None, + "parent_trace_context": { + "parent_workflow_run_id": "outermost-workflow-run", + "parent_node_execution_id": "outermost-node-execution", + "parent_app_id": None, + "parent_external_trace_id": "caller-trace", + "parent_trace_context": None, + }, } @@ -414,8 +431,9 @@ def test_workflow_tool_keeps_user_inputs_named_like_trace_runtime_keys( ), ] tool.set_parent_trace_context( - parent_workflow_run_id="outer-workflow-run-1", - parent_node_execution_id="outer-node-execution-1", + ParentTraceContext( + parent_workflow_run_id="outer-workflow-run-1", parent_node_execution_id="outer-node-execution-1" + ) ) monkeypatch.setattr(tool, "_get_app", lambda *args, **kwargs: None) @@ -441,10 +459,9 @@ def test_workflow_tool_keeps_user_inputs_named_like_trace_runtime_keys( call_kwargs = generate_calls[-1] assert call_kwargs["args"]["inputs"]["outer_workflow_run_id"] == "user-workflow-input" assert call_kwargs["args"]["inputs"]["outer_node_execution_id"] == "user-node-input" - assert call_kwargs["args"]["parent_trace_context"].model_dump() == { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - } + assert call_kwargs["args"]["parent_trace_context"] == ParentTraceContext( + parent_workflow_run_id="outer-workflow-run-1", parent_node_execution_id="outer-node-execution-1" + ) def test_workflow_tool_can_clear_parent_trace_context( @@ -454,8 +471,9 @@ def test_workflow_tool_can_clear_parent_trace_context( """Ensure reused WorkflowTool instances do not keep stale parent trace context.""" tool = _build_tool() tool.set_parent_trace_context( - parent_workflow_run_id="outer-workflow-run-1", - parent_node_execution_id="outer-node-execution-1", + ParentTraceContext( + parent_workflow_run_id="outer-workflow-run-1", parent_node_execution_id="outer-node-execution-1" + ) ) tool.clear_parent_trace_context() 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 e541b765b8df20..a344f31a4c9754 100644 --- a/api/tests/unit_tests/core/workflow/test_node_runtime.py +++ b/api/tests/unit_tests/core/workflow/test_node_runtime.py @@ -21,6 +21,7 @@ from core.app.llm.model_access import DifyCredentialsProvider, DifyModelFactory from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler from core.helper.ssrf_proxy import graphon_ssrf_proxy +from core.helper.trace_id_helper import ParentTraceContext from core.llm_generator.output_parser.errors import OutputParserError from core.llm_generator.output_parser.structured_output import _parse_structured_output from core.model_manager import QuotaManagedModelInstance @@ -1142,7 +1143,7 @@ def test_get_runtime_converts_graph_provider_type_for_tool_manager(runtime: Dify assert workflow_tool.provider_type == CoreToolProviderType.BUILT_IN @staticmethod - def test_get_runtime_stores_parent_trace_context_for_workflow_tools( + def test_get_runtime_chains_parent_trace_context_for_workflow_tools( runtime: DifyToolNodeRuntime, ) -> None: variable_pool: VariablePool = build_test_variable_pool( @@ -1153,6 +1154,11 @@ def test_get_runtime_stores_parent_trace_context_for_workflow_tools( ) workflow_runtime = MagicMock() workflow_runtime.runtime.runtime_parameters = {} + runtime._run_context.parent_trace_context = ParentTraceContext( + parent_workflow_run_id="outer-run-id", + parent_node_execution_id="outer-node-execution-id", + parent_external_trace_id="caller-trace", + ) node_data = ToolNodeData.model_validate( { "type": "tool", @@ -1175,14 +1181,20 @@ def test_get_runtime_stores_parent_trace_context_for_workflow_tools( node_execution_id="node-execution-id", ) - assert tool_runtime.raw.parent_trace_context.model_dump() == { - "parent_workflow_run_id": "workflow-run-id", - "parent_node_execution_id": "node-execution-id", - } + assert tool_runtime.raw.parent_trace_context == ParentTraceContext( + parent_workflow_run_id="workflow-run-id", + parent_node_execution_id="node-execution-id", + parent_app_id="app-id", + parent_trace_context=ParentTraceContext( + parent_workflow_run_id="outer-run-id", + parent_node_execution_id="outer-node-execution-id", + parent_external_trace_id="caller-trace", + ), + ) assert workflow_runtime.runtime.runtime_parameters == {} @staticmethod - def test_get_runtime_stores_trace_session_id_for_workflow_tools( + def test_get_runtime_passes_trace_session_and_caller_trace_to_workflow_tools( runtime: DifyToolNodeRuntime, ) -> None: variable_pool: VariablePool = build_test_variable_pool( @@ -1194,6 +1206,7 @@ def test_get_runtime_stores_trace_session_id_for_workflow_tools( workflow_runtime = MagicMock() workflow_runtime.runtime.runtime_parameters = {} runtime._run_context.trace_session_id = "session-1" + runtime._run_context.external_trace_id = "caller-trace" node_data = ToolNodeData.model_validate( { "type": "tool", @@ -1217,6 +1230,12 @@ def test_get_runtime_stores_trace_session_id_for_workflow_tools( ) assert tool_runtime.raw.trace_session_id == "session-1" + assert tool_runtime.raw.parent_trace_context == ParentTraceContext( + parent_workflow_run_id="workflow-run-id", + parent_node_execution_id="node-execution-id", + parent_app_id="app-id", + parent_external_trace_id="caller-trace", + ) assert workflow_runtime.runtime.runtime_parameters == {} @staticmethod @@ -2174,10 +2193,11 @@ def test_dify_tool_node_runtime_injects_outer_workflow_run_id_for_workflow_tools ) assert handle.raw.tool is runtime_tool - assert handle.raw.parent_trace_context.model_dump() == { - "parent_workflow_run_id": "outer-workflow-run-id", - "parent_node_execution_id": "node-execution-id", - } + assert handle.raw.parent_trace_context == ParentTraceContext( + parent_workflow_run_id="outer-workflow-run-id", + parent_node_execution_id="node-execution-id", + parent_app_id="app-id", + ) assert runtime_tool.runtime.runtime_parameters == {} get_runtime.assert_called_once() diff --git a/api/tests/unit_tests/domain/tracing/test_nested_workflow_trace.py b/api/tests/unit_tests/domain/tracing/test_nested_workflow_trace.py new file mode 100644 index 00000000000000..4dd2e804e639bb --- /dev/null +++ b/api/tests/unit_tests/domain/tracing/test_nested_workflow_trace.py @@ -0,0 +1,108 @@ +from collections.abc import Iterator +from datetime import UTC, datetime, timedelta + +from dify_trace_sdk.ids import span_id_for, trace_id_for +from dify_trace_sdk.model import Span, SpanKind, SpanRef, SpanStatus, Trace + +from domain.tracing.nested_workflow_trace import OuterTrace, TracePlacement, nested_trace_placement + +STARTED_AT = datetime(2026, 10, 1, 12, 0, 0, tzinfo=UTC) +TOOL_NODE = SpanRef(span_id=span_id_for("outer-tool-node"), start_time=STARTED_AT + timedelta(seconds=2)) + + +def outer(run: str, *, same_destination: bool = True, collected: bool = False) -> OuterTrace: + return OuterTrace( + trace_id=trace_id_for(run), + external_trace_id=f"{run}-caller", + root=SpanRef(span_id=span_id_for(run), start_time=STARTED_AT), + same_destination=same_destination, + collected=collected, + ) + + +def joined(run: str) -> TracePlacement: + return TracePlacement( + trace_id=trace_id_for(run), + external_trace_id=f"{run}-caller", + parent=TOOL_NODE, + trace_root=SpanRef(span_id=span_id_for(run), start_time=STARTED_AT), + ) + + +def span(span_id: str, parent_span_id: str | None) -> Span: + return Span( + span_id=span_id, + parent_span_id=parent_span_id, + name="workflow", + kind=SpanKind.CHAIN, + start_time=STARTED_AT + timedelta(seconds=3), + end_time=STARTED_AT + timedelta(seconds=4), + status=SpanStatus.OK, + ) + + +class TestNestedTracePlacement: + @staticmethod + def test_nested_run_hangs_from_the_tool_node_in_the_trace_of_its_outer_run() -> None: + assert nested_trace_placement([outer("parent")], tool_node=TOOL_NODE) == joined("parent") + + @staticmethod + def test_outer_runs_sharing_the_destination_join_the_outermost_trace() -> None: + placement = nested_trace_placement([outer("parent"), outer("grandparent")], tool_node=TOOL_NODE) + + assert placement == joined("grandparent") + + @staticmethod + def test_an_outer_run_exporting_elsewhere_starts_the_trace_below_it() -> None: + assert nested_trace_placement( + [outer("parent"), outer("grandparent", same_destination=False)], tool_node=TOOL_NODE + ) == joined("parent") + assert ( + nested_trace_placement([outer("parent", same_destination=False), outer("grandparent")], tool_node=TOOL_NODE) + is None + ) + + @staticmethod + def test_a_collected_outer_run_is_the_root_of_its_trace() -> None: + placement = nested_trace_placement([outer("parent", collected=True), outer("grandparent")], tool_node=TOOL_NODE) + + assert placement == joined("parent") + + @staticmethod + def test_a_nested_run_without_a_known_outer_trace_starts_its_own() -> None: + assert nested_trace_placement([], tool_node=TOOL_NODE) is None + + @staticmethod + def test_outer_runs_are_read_only_as_far_as_the_decision_needs() -> None: + def outer_traces() -> Iterator[OuterTrace]: + yield outer("parent", same_destination=False) + raise AssertionError("the run above an outer run exporting elsewhere was read") + + assert nested_trace_placement(outer_traces(), tool_node=TOOL_NODE) is None + + +class TestTracePlacement: + @staticmethod + def test_placed_trace_hangs_its_root_from_the_tool_node_of_the_outer_trace() -> None: + root_span_id, node_span_id = span_id_for("nested-run"), span_id_for("nested-node") + trace = Trace( + trace_id=trace_id_for("nested-run"), + external_trace_id=None, + session_id="session", + user_id="account", + parent=None, + trace_root=None, + spans=(span(root_span_id, None), span(node_span_id, root_span_id)), + ) + + placed = joined("parent").place(trace) + + assert placed.trace_id == trace_id_for("parent") + assert placed.external_trace_id == "parent-caller" + assert (placed.session_id, placed.user_id) == ("session", "account") + assert placed.parent == TOOL_NODE + assert placed.trace_root == SpanRef(span_id=span_id_for("parent"), start_time=STARTED_AT) + assert [(item.span_id, item.parent_span_id) for item in placed.spans] == [ + (root_span_id, TOOL_NODE.span_id), + (node_span_id, root_span_id), + ] diff --git a/api/tests/unit_tests/domain/tracing/test_normalization.py b/api/tests/unit_tests/domain/tracing/test_normalization.py index fa74baab3ad8a4..b2477593821ea3 100644 --- a/api/tests/unit_tests/domain/tracing/test_normalization.py +++ b/api/tests/unit_tests/domain/tracing/test_normalization.py @@ -179,6 +179,7 @@ def trace_of(*spans: Span) -> Trace: session_id="session", user_id="user", parent=None, + trace_root=None, spans=spans, ) @@ -272,6 +273,7 @@ def test_managed_external_parent_is_preserved_as_explicit_link_when_content_off( session_id="session", user_id=None, parent=None, + trace_root=None, spans=(root,), ) remapped = remap_outer_trace(outer, trace_id="1" * 32, capture_content=False, synthetic_span_ids=frozenset()) 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 index 8fe1b1d90b659c..6775db0123b354 100644 --- a/api/tests/unit_tests/domain/tracing/test_trace_run_agg.py +++ b/api/tests/unit_tests/domain/tracing/test_trace_run_agg.py @@ -95,6 +95,7 @@ def business_terminal(run: TraceRunAgg) -> Trace: session_id="conversation", user_id="user", parent=None, + trace_root=None, spans=( Span( span_id=span_id_for(run.source_id), 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 90502db2cca719..05dccbf8854690 100644 --- a/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py +++ b/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py @@ -334,9 +334,7 @@ def create_scope( 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ) def emit(self, event: TraceEvent) -> None: 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 5aa4c413cefc6e..b730c487ff876b 100644 --- a/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py +++ b/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py @@ -98,6 +98,7 @@ def make_workflow_info(**overrides) -> WorkflowTraceInfo: "completion_tokens": 40, "file_list": [], "query": "hello", + "parent": None, "start_time": _T0, "end_time": _T1, "metadata": { diff --git a/api/tests/unit_tests/extensions/test_tracing.py b/api/tests/unit_tests/extensions/test_tracing.py index 64ee89e7f73309..446ac81843a062 100644 --- a/api/tests/unit_tests/extensions/test_tracing.py +++ b/api/tests/unit_tests/extensions/test_tracing.py @@ -2,7 +2,7 @@ from dify_agent.layers.execution_context import ApplicationTracingContext from dify_trace_sdk.emitter import TraceScope -from dify_trace_sdk.events import TraceCorrelation, TraceEvent +from dify_trace_sdk.events import TraceCorrelation, TraceEvent, TraceParent from flask import Flask from core.helper.trace_id_helper import ParentTraceContext @@ -113,7 +113,17 @@ def test_correlation_reads_the_caller_extras() -> None: { "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"}, + "parent_trace_context": { + "parent_workflow_run_id": "outer-run", + "parent_node_execution_id": "outer-node", + "parent_app_id": "outer-app", + "parent_trace_context": { + "parent_workflow_run_id": "outermost-run", + "parent_node_execution_id": "outermost-node", + "parent_app_id": "outermost-app", + "parent_external_trace_id": "outermost-caller-trace", + }, + }, } ) typed = correlation_from_extras( @@ -123,12 +133,25 @@ def test_correlation_reads_the_caller_extras() -> 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", + parent=TraceParent( + workflow_run_id="outer-run", + app_id="outer-app", + node_execution_id="outer-node", + external_trace_id=None, + parent=TraceParent( + workflow_run_id="outermost-run", + app_id="outermost-app", + node_execution_id="outermost-node", + external_trace_id="outermost-caller-trace", + parent=None, + ), + ), + ) + assert typed.parent == TraceParent( + workflow_run_id="outer-run", app_id=None, node_execution_id=None, external_trace_id=None, parent=None ) - 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 + external_trace_id=None, session_id=None, parent=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 2cc696c336ba2a..a8345389423d20 100644 --- a/api/tests/unit_tests/repositories/test_application_trace_repository.py +++ b/api/tests/unit_tests/repositories/test_application_trace_repository.py @@ -36,7 +36,7 @@ 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.delivery import CollectedTrace, StoredTraceDestination, TraceRecordSource from services.tracing.errors import ( TraceCollectionConflictError, TraceCollectionError, @@ -224,6 +224,7 @@ def agent_turn(run: TraceRun, *, root_span_id: str) -> Trace: session_id=run.session_id, user_id=None, parent=None, + trace_root=None, spans=(root,), ) @@ -332,6 +333,7 @@ def workflow_info(row: WorkflowRun, parent: WorkflowRun | None) -> WorkflowTrace file_list=[], invoked_by=None, query="", + parent=None, metadata=metadata, start_time=row.created_at, end_time=row.finished_at, @@ -721,21 +723,22 @@ def test_auxiliary_content_is_suppressed_for_logstore_chatflow(stored: StoredSou @staticmethod @pytest.mark.parametrize("stored", ["workflow"], indirect=True) - @pytest.mark.parametrize("mismatch", [None, "tenant_id", "app_id"]) - def test_linked_parent_requires_logstore_run_and_durable_owner_to_agree( - stored: StoredSource, mismatch: str | None - ) -> None: + def test_collected_trace_is_the_frozen_agent_trace_of_the_runs_own_app(stored: StoredSource) -> None: parent = stored.start() - info = workflow_info(_dict_to_workflow_run(stored.row), _dict_to_workflow_run(stored.row)) - if mismatch: - stored.row[mismatch] = identifier("foreign") - linked = stored.service().link_managed_parent(tenant_id=stored.case.run.tenant_id, info=info) - if mismatch: - assert linked is info - else: - 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 + service = stored.service() + + collected = service.collected_trace( + tenant_id=parent.tenant_id, app_id=parent.app_id, workflow_run_id=parent.source_id + ) + foreign = service.collected_trace( + tenant_id=parent.tenant_id, app_id=identifier("foreign"), workflow_run_id=parent.source_id + ) + + assert collected == CollectedTrace( + trace_id=parent.trace_id, + destination=StoredTraceDestination(provider_id=parent.provider, config=parent.provider_config), + ) + assert foreign == CollectedTrace(trace_id=None, destination=None) @staticmethod def test_sql_workflow_source_still_uses_scoped_configured_reader( diff --git a/api/tests/unit_tests/repositories/test_trace_record_repository.py b/api/tests/unit_tests/repositories/test_trace_record_repository.py index 51f882d3fdaa62..8b2957575ee27b 100644 --- a/api/tests/unit_tests/repositories/test_trace_record_repository.py +++ b/api/tests/unit_tests/repositories/test_trace_record_repository.py @@ -13,9 +13,24 @@ from sqlalchemy import create_engine from sqlalchemy.orm import Session, sessionmaker -from models.enums import ConversationFromSource, ConversationStatus, EndUserType, InvokeFrom +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.trace_record_repository import TraceRecordRepository +from services.tracing.delivery import OuterRunRoot from services.tracing.errors import SourceNotFoundError, SourceNotReadyError, TracingStoreUnavailableError from services.tracing.records import MessageTraceInfo, ModerationTraceInfo, ToolTraceInfo, TraceMessage @@ -23,6 +38,10 @@ 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" +CHATFLOW_RUN_ID = "outer-chatflow-run" +TOOL_NODE_ID = "outer-tool-node" @pytest.fixture @@ -91,6 +110,77 @@ def records(sqlite_session: Session) -> None: 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 running workflow run whose tool node starts nested runs, and the chatflow run the message answers.""" + 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) + chat_app = sqlite_session.get(App, APP_ID) + message = sqlite_session.get(Message, MESSAGE_ID) + assert chat_app is not None + assert message is not None + chat_app.mode = AppMode.ADVANCED_CHAT + message.workflow_run_id = CHATFLOW_RUN_ID + sqlite_session.add_all( + [workflow_app, _workflow_run(OUTER_RUN_ID, WORKFLOW_APP_ID), tool_node, _workflow_run(CHATFLOW_RUN_ID, APP_ID)] + ) + sqlite_session.commit() + + def _assembler(sessions: sessionmaker[Session]) -> TraceRecordRepository: return TraceRecordRepository(sessions, files_url="https://files.example.com", enterprise_details=False) @@ -105,8 +195,7 @@ def _event(payload: TracePayload, *, tenant_id: str = TENANT_ID) -> TraceEvent: correlation=TraceCorrelation( external_trace_id="caller-trace", session_id="session-1", - parent_workflow_run_id=None, - parent_node_execution_id=None, + parent=None, ), payload=payload, ) @@ -235,3 +324,51 @@ def test_an_unreachable_database_is_reported_as_unavailable() -> None: with pytest.raises(TracingStoreUnavailableError): _assembler(sessions).assemble(_event(MessageCompleted(message_id=MESSAGE_ID, conversation_id="conversation-1"))) + + +@pytest.mark.usefixtures("records", "outer_runs") +def test_a_workflow_trace_starts_with_its_run_and_a_chatflow_trace_with_its_message( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = _assembler(sqlite_session_factory) + + assert repository.root(tenant_id=TENANT_ID, app_id=WORKFLOW_APP_ID, workflow_run_id=OUTER_RUN_ID) == OuterRunRoot( + root_id=OUTER_RUN_ID, started_at=CREATED_AT + timedelta(seconds=1) + ) + assert repository.root(tenant_id=TENANT_ID, app_id=APP_ID, workflow_run_id=CHATFLOW_RUN_ID) == OuterRunRoot( + root_id=MESSAGE_ID, started_at=CREATED_AT + ) + + +@pytest.mark.usefixtures("records", "outer_runs") +def test_outer_runs_and_their_tool_nodes_are_read_through_their_owners( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = _assembler(sqlite_session_factory) + + def started_at(tenant_id: str, app_id: str, workflow_run_id: str) -> datetime | None: + return repository.node_started_at( + tenant_id=tenant_id, app_id=app_id, workflow_run_id=workflow_run_id, node_execution_id=TOOL_NODE_ID + ) + + assert started_at(TENANT_ID, WORKFLOW_APP_ID, OUTER_RUN_ID) == CREATED_AT + timedelta(seconds=2) + assert started_at("other-tenant", WORKFLOW_APP_ID, OUTER_RUN_ID) is None + assert started_at(TENANT_ID, APP_ID, OUTER_RUN_ID) is None + assert started_at(TENANT_ID, WORKFLOW_APP_ID, CHATFLOW_RUN_ID) is None + assert repository.root(tenant_id="other-tenant", app_id=WORKFLOW_APP_ID, workflow_run_id=OUTER_RUN_ID) is None + assert repository.root(tenant_id=TENANT_ID, app_id=APP_ID, workflow_run_id=OUTER_RUN_ID) is None + + +@pytest.mark.usefixtures("records", "outer_runs") +def test_a_chatflow_run_has_no_root_until_its_message_names_it( + sqlite_session_factory: sessionmaker[Session], sqlite_session: Session +) -> None: + message = sqlite_session.get(Message, MESSAGE_ID) + assert message is not None + message.workflow_run_id = None + sqlite_session.commit() + + assert ( + _assembler(sqlite_session_factory).root(tenant_id=TENANT_ID, app_id=APP_ID, workflow_run_id=CHATFLOW_RUN_ID) + is None + ) diff --git a/api/tests/unit_tests/services/test_workflow_persistence.py b/api/tests/unit_tests/services/test_workflow_persistence.py index ea38df25573781..ad3f839319fc10 100644 --- a/api/tests/unit_tests/services/test_workflow_persistence.py +++ b/api/tests/unit_tests/services/test_workflow_persistence.py @@ -8,7 +8,7 @@ import pytest from dify_trace_sdk.emitter import TraceScope -from dify_trace_sdk.events import TraceCorrelation, TraceEvent, WorkflowRunCompleted +from dify_trace_sdk.events import TraceCorrelation, TraceEvent, TraceParent, WorkflowRunCompleted from flask import Flask from sqlalchemy import event as sqlalchemy_event from sqlalchemy import select @@ -293,9 +293,7 @@ def _trace_scope(recorder: _TraceRecorder, correlation: TraceCorrelation) -> Tra def _empty_correlation() -> TraceCorrelation: - return TraceCorrelation( - external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None - ) + return TraceCorrelation(external_trace_id=None, session_id=None, parent=None) def _naive_utc_now() -> datetime: @@ -1104,8 +1102,13 @@ def test_handle_graph_run_succeeded_emits_run_tokens_with_scope_correlation( 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", + parent=TraceParent( + workflow_run_id="outer-workflow-run-1", + app_id="outer-app-1", + node_execution_id="outer-node-execution-1", + external_trace_id=None, + parent=None, + ), ) layer, _, _, runtime_state = _make_layer( system_variables, diff --git a/api/tests/unit_tests/services/tracing/test_adapters.py b/api/tests/unit_tests/services/tracing/test_adapters.py index 5afba1f01b2a70..ff46c3a61805cb 100644 --- a/api/tests/unit_tests/services/tracing/test_adapters.py +++ b/api/tests/unit_tests/services/tracing/test_adapters.py @@ -24,9 +24,7 @@ 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 -) +CORRELATION = TraceCorrelation(external_trace_id=None, session_id=None, parent=None) class _Broker: @@ -152,6 +150,7 @@ def test_export_retries_persist_the_trace_and_run_after_their_delay(broker: _Bro session_id=None, user_id=None, parent=None, + trace_root=None, spans=( Span( span_id="b" * 16, diff --git a/api/tests/unit_tests/services/tracing/test_canonical.py b/api/tests/unit_tests/services/tracing/test_canonical.py index 1c22b6f305b150..de91f7b270d76a 100644 --- a/api/tests/unit_tests/services/tracing/test_canonical.py +++ b/api/tests/unit_tests/services/tracing/test_canonical.py @@ -6,7 +6,7 @@ from graphon.entities import WorkflowNodeExecution from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus -from services.tracing.canonical import CanonicalTraceBuilder, message_trace_id +from services.tracing.canonical import CanonicalTraceBuilder, root_trace_id from services.tracing.records import ( MessageTraceInfo, PromptGenerationTraceInfo, @@ -116,6 +116,7 @@ def _workflow_record(*, message_id: str | None, metadata: dict[str, object]) -> file_list=[], invoked_by=None, query="hi" if message_id else "", + parent=None, start_time=STARTED_AT, end_time=STARTED_AT + timedelta(seconds=3), metadata=metadata, @@ -144,7 +145,7 @@ 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.trace_id == root_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 diff --git a/api/tests/unit_tests/services/tracing/test_delivery.py b/api/tests/unit_tests/services/tracing/test_delivery.py index 3b80818652068d..d33f3cce2994ae 100644 --- a/api/tests/unit_tests/services/tracing/test_delivery.py +++ b/api/tests/unit_tests/services/tracing/test_delivery.py @@ -6,16 +6,20 @@ 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.events import MessageCompleted, TraceCorrelation, TraceEvent, TraceParent, WorkflowRunCompleted +from dify_trace_sdk.ids import span_id_for, trace_id_for +from dify_trace_sdk.model import SpanRef, 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 ( + CollectedTrace, ExportJob, ExportRetryQueue, + OuterRunQuery, + OuterRunRoot, StoredTraceDestination, TraceCollector, TraceDeliveryService, @@ -23,7 +27,7 @@ TraceRecordSource, retry_delay_seconds, ) -from services.tracing.errors import SourceNotReadyError, TraceCollectorUnavailableError +from services.tracing.errors import SourceNotReadyError, TraceCollectorUnavailableError, TracingStoreUnavailableError 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 @@ -32,6 +36,9 @@ MESSAGE_ID = "6f1c0b52-6ad4-4c4e-9d7c-2a3f0d0c7a11" TENANT_ID = "tenant-1" APP_ID = "app-1" +NESTED_RUN_ID = "nested-run" +TOOL_STARTED_AT = STARTED_AT + timedelta(seconds=2) +SAME_DESTINATION = StoredTraceDestination(provider_id="langfuse", config={"endpoint": "https://collector.example"}) class _Config(ProviderConfig): @@ -92,6 +99,33 @@ def active_destination(self, *, tenant_id: str, app_id: str) -> StoredTraceDesti return self.destinations.get((tenant_id, app_id)) +class _OuterRuns(OuterRunQuery): + """Outer run roots by (tenant, app, run) and node starts by (tenant, app, run, node).""" + + def __init__(self) -> None: + self.roots: dict[tuple[str, str, str], OuterRunRoot] = {} + self.node_starts: dict[tuple[str, str, str, str], datetime] = {} + self.unreadable = False + + @override + def root(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> OuterRunRoot | None: + if self.unreadable: + raise TracingStoreUnavailableError + return self.roots.get((tenant_id, app_id, workflow_run_id)) + + @override + def node_started_at( + self, *, tenant_id: str, app_id: str, workflow_run_id: str, node_execution_id: str + ) -> datetime | None: + return self.node_starts.get((tenant_id, app_id, workflow_run_id, node_execution_id)) + + def add(self, link: TraceParent, root: OuterRunRoot) -> None: + assert link.app_id is not None + assert link.node_execution_id is not None + self.roots[(TENANT_ID, link.app_id, link.workflow_run_id)] = root + self.node_starts[(TENANT_ID, link.app_id, link.workflow_run_id, link.node_execution_id)] = TOOL_STARTED_AT + + class _Retries(ExportRetryQueue): def __init__(self) -> None: self.scheduled: list[tuple[ExportJob, int]] = [] @@ -106,6 +140,7 @@ def __init__(self, *, captures: bool, fails: bool) -> None: self.captures = captures self.fails = fails self.seen: list[tuple[str, BaseTraceInfo]] = [] + self.collected: dict[str, CollectedTrace] = {} @override def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: @@ -115,8 +150,8 @@ def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: return self.captures @override - def link_managed_parent(self, *, tenant_id: str, info: WorkflowTraceInfo) -> WorkflowTraceInfo: - return info + def collected_trace(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> CollectedTrace | None: + return self.collected.get(workflow_run_id) @dataclass @@ -125,11 +160,13 @@ class _Delivery: egress: RecordingEgress destinations: _Destinations retries: _Retries = field(default_factory=_Retries) + outer_runs: _OuterRuns = field(default_factory=_OuterRuns) enterprise: list[BaseTraceInfo] = field(default_factory=list) def service(self, *, enterprise: bool, collector: _Collector | None) -> TraceDeliveryService: return TraceDeliveryService( records=self.records, + outer_runs=self.outer_runs, destinations=self.destinations, providers=InstalledTraceProviders(_Provider), codec=ProviderConfigCodec(WorkspaceBoundCipher()), @@ -152,9 +189,7 @@ def _event() -> TraceEvent: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), payload=MessageCompleted(message_id=MESSAGE_ID, conversation_id="conversation-1"), ) @@ -195,12 +230,71 @@ def _record(*, workflow_run_id: str | None) -> MessageTraceInfo: ) +def _workflow_event() -> TraceEvent: + return _event().model_copy( + update={"payload": WorkflowRunCompleted(workflow_run_id=NESTED_RUN_ID, conversation_id=None, total_tokens=None)} + ) + + +def _nested_record(parent: TraceParent) -> WorkflowTraceInfo: + return WorkflowTraceInfo( + node_executions=[], + conversation_id=None, + message_started_at=None, + workflow_app_log_id=None, + workflow_id="workflow-1", + tenant_id=TENANT_ID, + workflow_run_id=NESTED_RUN_ID, + workflow_run_elapsed_time=1, + 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="", + parent=parent, + start_time=STARTED_AT + timedelta(seconds=3), + end_time=STARTED_AT + timedelta(seconds=4), + metadata={}, + ) + + +def _link(run_id: str, *, external_trace_id: str | None, parent: TraceParent | None) -> TraceParent: + return TraceParent( + workflow_run_id=run_id, + app_id=f"{run_id}-app", + node_execution_id=f"{run_id}-tool-node", + external_trace_id=external_trace_id, + parent=parent, + ) + + +def _nested_delivery(parent: TraceParent) -> _Delivery: + """A nested run whose outer runs are all persisted and trace where it does.""" + delivery = _delivery(_nested_record(parent), [response(200, {})]) + link: TraceParent | None = parent + while link is not None: + assert link.app_id is not None + delivery.destinations.destinations[(TENANT_ID, link.app_id)] = SAME_DESTINATION + delivery.outer_runs.add(link, OuterRunRoot(root_id=link.workflow_run_id, started_at=STARTED_AT)) + link = link.parent + return delivery + + +def _tool_node(run_id: str) -> SpanRef: + return SpanRef(span_id=span_id_for(f"{run_id}-tool-node"), start_time=TOOL_STARTED_AT) + + 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}), + destinations=_Destinations({(TENANT_ID, APP_ID): SAME_DESTINATION}), ) @@ -365,3 +459,82 @@ def test_destinations_that_cannot_be_used_send_nothing(stored: StoredTraceDestin delivery.service(enterprise=False, collector=None).deliver(_event(), enterprise_delivered=False) assert delivery.sent() == [] + + +def test_nested_run_hangs_from_the_tool_node_in_the_trace_of_its_outer_run() -> None: + delivery = _nested_delivery(_link("outer-run", external_trace_id=None, parent=None)) + + delivery.service(enterprise=False, collector=None).deliver(_workflow_event(), enterprise_delivered=False) + + (trace,) = delivery.sent() + assert trace.trace_id == trace_id_for("outer-run") + assert trace.external_trace_id is None + assert trace.parent == _tool_node("outer-run") + assert trace.trace_root == SpanRef(span_id=span_id_for("outer-run"), start_time=STARTED_AT) + assert trace.root.span_id == span_id_for(NESTED_RUN_ID) + assert trace.root.parent_span_id == span_id_for("outer-run-tool-node") + + +def test_outer_runs_sharing_the_destination_join_the_trace_of_the_outermost_one() -> None: + chatflow = _link("chatflow-run", external_trace_id="caller-trace", parent=None) + delivery = _nested_delivery(_link("outer-run", external_trace_id=None, parent=chatflow)) + message_started_at = STARTED_AT - timedelta(seconds=1) + delivery.outer_runs.roots[(TENANT_ID, "chatflow-run-app", "chatflow-run")] = OuterRunRoot( + root_id="chatflow-message", started_at=message_started_at + ) + + delivery.service(enterprise=False, collector=None).deliver(_workflow_event(), enterprise_delivered=False) + + (trace,) = delivery.sent() + assert (trace.trace_id, trace.external_trace_id) == (trace_id_for("caller-trace"), "caller-trace") + assert trace.parent == _tool_node("outer-run") + assert trace.trace_root == SpanRef(span_id=span_id_for("chatflow-message"), start_time=message_started_at) + + +def test_a_nested_run_starts_its_own_trace_below_an_outer_run_tracing_elsewhere() -> None: + delivery = _nested_delivery(_link("outer-run", external_trace_id=None, parent=None)) + delivery.destinations.destinations[(TENANT_ID, "outer-run-app")] = StoredTraceDestination( + provider_id="langfuse", config={"endpoint": "https://elsewhere.example"} + ) + + delivery.service(enterprise=False, collector=None).deliver(_workflow_event(), enterprise_delivered=False) + + (trace,) = delivery.sent() + assert (trace.trace_id, trace.parent, trace.trace_root) == (trace_id_for(NESTED_RUN_ID), None, None) + + +def test_a_collected_outer_run_lends_its_agent_trace_once_collection_starts() -> None: + collector = _Collector(captures=False, fails=False) + agent_trace = CollectedTrace(trace_id=trace_id_for("agent-trace"), destination=SAME_DESTINATION) + collector.collected["outer-run"] = agent_trace + delivery = _nested_delivery(_link("outer-run", external_trace_id="caller-trace", parent=None)) + + delivery.service(enterprise=False, collector=collector).deliver(_workflow_event(), enterprise_delivered=False) + + (trace,) = delivery.sent() + assert (trace.trace_id, trace.external_trace_id) == (trace_id_for("agent-trace"), None) + assert trace.parent == _tool_node("outer-run") + + +def test_a_collected_outer_run_lends_no_trace_before_collection_starts() -> None: + collector = _Collector(captures=False, fails=False) + collector.collected["outer-run"] = CollectedTrace(trace_id=None, destination=None) + delivery = _nested_delivery(_link("outer-run", external_trace_id=None, parent=None)) + + delivery.service(enterprise=False, collector=collector).deliver(_workflow_event(), enterprise_delivered=False) + + (trace,) = delivery.sent() + assert (trace.trace_id, trace.parent) == (trace_id_for(NESTED_RUN_ID), None) + + +def test_a_nested_run_starts_its_own_trace_while_its_outer_records_are_unavailable() -> None: + unpersisted = _nested_delivery(_link("outer-run", external_trace_id=None, parent=None)) + unpersisted.outer_runs.node_starts.clear() + unreadable = _nested_delivery(_link("outer-run", external_trace_id=None, parent=None)) + unreadable.outer_runs.unreadable = True + + for delivery in (unpersisted, unreadable): + delivery.service(enterprise=False, collector=None).deliver(_workflow_event(), enterprise_delivered=False) + + (trace,) = delivery.sent() + assert (trace.trace_id, trace.parent) == (trace_id_for(NESTED_RUN_ID), None) diff --git a/api/tests/unit_tests/services/tracing/test_recovery.py b/api/tests/unit_tests/services/tracing/test_recovery.py index 40c5225030c0dd..f8db25d6df9163 100644 --- a/api/tests/unit_tests/services/tracing/test_recovery.py +++ b/api/tests/unit_tests/services/tracing/test_recovery.py @@ -272,9 +272,7 @@ def trace_event(run: TraceRun, payload: TracePayload) -> TraceEvent: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), payload=payload, ) @@ -552,9 +550,7 @@ def send(message: Message, *, application_generate_entity: ChatAppGenerateEntity 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), ), ) stop = MessageQueueMessage( 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 ccc648d5d37fee..3104feba413be4 100644 --- a/api/tests/unit_tests/services/tracing/test_runtime_bridge.py +++ b/api/tests/unit_tests/services/tracing/test_runtime_bridge.py @@ -18,10 +18,11 @@ ModerationChecked, TraceCorrelation, TraceEvent, + TraceParent, TracePayload, WorkflowRunCompleted, ) -from dify_trace_sdk.ids import span_id_for +from dify_trace_sdk.ids import span_id_for, span_uuid, trace_id_for, trace_uuid from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace from dify_trace_sdk.testing import RecordingEgress, response from flask import Flask @@ -310,9 +311,7 @@ def trace_event(run: TraceRun, payload: TracePayload) -> TraceEvent: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), payload=payload, ) @@ -349,6 +348,7 @@ def terminal(collection: Collection) -> None: session_id=run.session_id, user_id=None, parent=None, + trace_root=None, spans=(root,), ) capture(collection.repository, run.id, trace) @@ -384,8 +384,10 @@ def setup(sessions: sessionmaker[Session], name: str, mode: AppMode, *, snapshot 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 + records = record_source(sessions) TraceDeliveryService( - records=record_source(sessions), + records=records, + outer_runs=records, destinations=SQLAlchemyAppTracingConfigRepository(session_factory=sessions), providers=InstalledTraceProviders(LangfuseProvider), codec=ProviderConfigCodec(WorkspaceBoundCipher()), @@ -420,8 +422,15 @@ def workflow_event(row: WorkflowRun, parent: WorkflowRun | None) -> TraceEvent: 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, + parent=TraceParent( + workflow_run_id=parent.id, + app_id=parent.app_id, + node_execution_id=identifier("parent-tool-node"), + external_trace_id=None, + parent=None, + ) + if parent + else None, ), payload=WorkflowRunCompleted(workflow_run_id=row.id, conversation_id=None, total_tokens=None), ) @@ -477,6 +486,46 @@ def workflow(case: Case, name: str, *, has_agent: bool, snapshot: bool) -> tuple return row, trace_run +def tool_node(case: Case, row: WorkflowRun) -> None: + """Persist the tool node execution of ``row`` that starts the nested runs of these cases.""" + with case.collection.sessions.begin() as session: + node = WorkflowNodeExecutionModel( + id=identifier("parent-tool-node"), + node_execution_id=identifier("parent-tool-node"), + tenant_id=row.tenant_id, + app_id=row.app_id, + workflow_id=row.workflow_id, + workflow_run_id=row.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=row.created_by, + created_by_role=CreatorUserRole.ACCOUNT, + ) + node.created_at = row.created_at + timedelta(seconds=1) + session.add(node) + + +def nested_root(destination: RecordingEgress, trace_id: str, child: WorkflowRun) -> dict[str, object]: + """The Langfuse observation of a nested run's root, after checking it was sent as a fragment.""" + events = ingested(destination) + assert all(item["type"] != "trace-create" for item in events) + bodies = [item["body"] for item in events if isinstance(item["body"], dict)] + (root,) = [body for body in bodies if body["id"] == str(span_uuid(trace_id, span_id_for(child.id)))] + assert root["traceId"] == str(trace_uuid(trace_id)) + assert root["parentObservationId"] == str(span_uuid(trace_id, span_id_for(identifier("parent-tool-node")))) + return root + + 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} @@ -505,6 +554,7 @@ def workflow_info(row: WorkflowRun, parent: WorkflowRun | None) -> WorkflowTrace file_list=[], invoked_by=None, query="", + parent=None, metadata=metadata, start_time=row.created_at, end_time=row.finished_at, @@ -628,11 +678,12 @@ def test_auxiliary_events_resolve_owner_without_bypassing_content_policy( assert run_count(case.collection.sessions) == 0 @staticmethod - def test_plain_nested_workflow_keeps_its_parent_context_for_the_provider( + def test_plain_nested_workflow_hangs_from_the_tool_node_of_its_parent( sqlite_session_factory: sessionmaker[Session], ) -> None: parent_case = setup(sqlite_session_factory, "plain-parent", AppMode.WORKFLOW, snapshot=False) parent, _ = workflow(parent_case, "plain-parent", has_agent=False, snapshot=False) + tool_node(parent_case, parent) child_case = setup(sqlite_session_factory, "plain-child", AppMode.WORKFLOW, snapshot=False) child, _ = workflow(child_case, "plain-child", has_agent=False, snapshot=False) add_app_owner(child_case.collection, "plain-child") @@ -640,15 +691,19 @@ def test_plain_nested_workflow_keeps_its_parent_context_for_the_provider( 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 + assert ( + child_case.service.collected_trace( + tenant_id=parent.tenant_id, app_id=parent.app_id, workflow_run_id=parent.id + ) + is None + ) 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"] + root = nested_root(destination, trace_id_for(parent.id), child) + metadata = root["metadata"] + assert isinstance(metadata, dict) + assert metadata["parent_workflow_run_id"] == parent.id + assert metadata["parent_node_execution_id"] == identifier("parent-tool-node") assert run_count(sqlite_session_factory) == 0 @staticmethod @@ -674,32 +729,40 @@ def test_plain_chatflow_auxiliary_events_keep_the_provider_route( assert len(destination.requests) == 1 @staticmethod - def test_provider_exported_child_of_managed_parent_names_the_collected_parent( + def test_provider_exported_child_of_a_collected_parent_hangs_in_its_agent_trace( sqlite_session_factory: sessionmaker[Session], ) -> None: 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 + tool_node(parent_case, parent) 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(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")) + collected = child_case.service.collected_trace( + tenant_id=parent.tenant_id, app_id=parent.app_id, workflow_run_id=parent.id + ) + assert collected is not None + assert collected.trace_id == parent_trace.trace_id + assert collected.destination is not None + assert collected.destination.provider_id == parent_trace.provider 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"] + nested_root(destination, parent_trace.trace_id, child) + + @staticmethod + def test_a_collected_parent_lends_no_trace_before_its_collection_starts( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + parent_case = setup(sqlite_session_factory, "pending-parent", AppMode.WORKFLOW, snapshot=False) + parent, _ = workflow(parent_case, "pending-parent", has_agent=True, snapshot=False) + + collected = parent_case.service.collected_trace( + tenant_id=parent.tenant_id, app_id=parent.app_id, workflow_run_id=parent.id + ) + + assert collected is not None + assert (collected.trace_id, collected.destination) == (None, None) @staticmethod def test_collector_database_failure_redelivers_without_provider_fallback( @@ -1414,6 +1477,7 @@ def test_hidden_child_skip_keeps_normal_root_agent_trace_assemblable( file_list=[], invoked_by=None, query="", + parent=None, metadata={"app_id": ids["root"]}, start_time=workflow.created_at, end_time=workflow.finished_at, diff --git a/api/tests/unit_tests/services/tracing/test_scopes.py b/api/tests/unit_tests/services/tracing/test_scopes.py index f269375ea3c8b0..ffd9c410991725 100644 --- a/api/tests/unit_tests/services/tracing/test_scopes.py +++ b/api/tests/unit_tests/services/tracing/test_scopes.py @@ -6,9 +6,7 @@ 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 -) +CORRELATION = TraceCorrelation(external_trace_id="caller-trace", session_id=None, parent=None) ENABLED = AppTraceSettings(enabled=True, tracing_provider="langfuse") diff --git a/api/tests/unit_tests/services/tracing/test_service.py b/api/tests/unit_tests/services/tracing/test_service.py index 1dfe6b4a4a3276..5a07dd21472964 100644 --- a/api/tests/unit_tests/services/tracing/test_service.py +++ b/api/tests/unit_tests/services/tracing/test_service.py @@ -354,9 +354,7 @@ def trace_event(case: _ApplicationTraceCollectionCollection, payload: TracePaylo 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), payload=payload, ) @@ -427,6 +425,7 @@ def _application_trace_collection_terminal(collection: _ApplicationTraceCollecti session_id=run.session_id, user_id=None, parent=None, + trace_root=None, spans=(root,), ) with collection.change() as current: 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 d02af7d055d76e..86a22362474a6b 100644 --- a/api/tests/unit_tests/tasks/test_application_trace_task.py +++ b/api/tests/unit_tests/tasks/test_application_trace_task.py @@ -128,9 +128,7 @@ def test_worker_records_carry_no_enterprise_names(sqlite_session_factory: sessio 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), payload=MessageCompleted(message_id=case.run.source_id, conversation_id=conversation_id), ) ) diff --git a/api/tests/unit_tests/tasks/test_trace_delivery_task.py b/api/tests/unit_tests/tasks/test_trace_delivery_task.py index 9e70a12ba7d04d..afb4ff04cd0c3a 100644 --- a/api/tests/unit_tests/tasks/test_trace_delivery_task.py +++ b/api/tests/unit_tests/tasks/test_trace_delivery_task.py @@ -15,8 +15,11 @@ from services.tracing.adapters import payload_key from services.tracing.canonical import CanonicalTraceBuilder from services.tracing.delivery import ( + CollectedTrace, ExportJob, ExportRetryQueue, + OuterRunQuery, + OuterRunRoot, StoredTraceDestination, TraceCollector, TraceDeliveryService, @@ -25,7 +28,7 @@ ) from services.tracing.errors import SourceNotFoundError, SourceNotReadyError from services.tracing.providers import ProviderConfigCodec -from services.tracing.records import BaseTraceInfo, GenerateNameTraceInfo, WorkflowTraceInfo +from services.tracing.records import BaseTraceInfo, GenerateNameTraceInfo 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 @@ -88,8 +91,20 @@ 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 + def collected_trace(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> CollectedTrace | None: + raise AssertionError("no nested workflow run is delivered") + + +class _NoOuterRuns(OuterRunQuery): + @override + def root(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> OuterRunRoot | None: + raise AssertionError("no nested workflow run is delivered") + + @override + def node_started_at( + self, *, tenant_id: str, app_id: str, workflow_run_id: str, node_execution_id: str + ) -> datetime | None: + raise AssertionError("no nested workflow run is delivered") class _Delivery: @@ -100,6 +115,7 @@ def __init__(self) -> None: def service(self, collector: TraceCollector | None) -> TraceDeliveryService: return TraceDeliveryService( records=self.records, + outer_runs=_NoOuterRuns(), destinations=self.destinations, providers=InstalledTraceProviders(), codec=ProviderConfigCodec(WorkspaceBoundCipher()), @@ -128,9 +144,7 @@ def _event() -> TraceEvent: 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 - ), + correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), payload=MessageCompleted(message_id="message-1", conversation_id=None), ) From 1cd9850bc5e3d544c7c41b19cffdffdf7e5c0043 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Thu, 29 Oct 2026 00:19:40 +0800 Subject: [PATCH 05/17] test(workflow): cover the parent trace context handed to workflow tools A workflow tool node passes its run's chained parent trace context to the tool it invokes. Exercising that hand-off keeps the KnowledgeFS changed-glue coverage gate above its minimum. --- .../core/workflow/test_node_runtime.py | 21 +++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) 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 a344f31a4c9754..34d3075ccde8b4 100644 --- a/api/tests/unit_tests/core/workflow/test_node_runtime.py +++ b/api/tests/unit_tests/core/workflow/test_node_runtime.py @@ -1143,7 +1143,7 @@ def test_get_runtime_converts_graph_provider_type_for_tool_manager(runtime: Dify assert workflow_tool.provider_type == CoreToolProviderType.BUILT_IN @staticmethod - def test_get_runtime_chains_parent_trace_context_for_workflow_tools( + def test_workflow_tools_receive_the_chained_parent_trace_context( runtime: DifyToolNodeRuntime, ) -> None: variable_pool: VariablePool = build_test_variable_pool( @@ -1181,7 +1181,7 @@ def test_get_runtime_chains_parent_trace_context_for_workflow_tools( node_execution_id="node-execution-id", ) - assert tool_runtime.raw.parent_trace_context == ParentTraceContext( + chained = ParentTraceContext( parent_workflow_run_id="workflow-run-id", parent_node_execution_id="node-execution-id", parent_app_id="app-id", @@ -1191,8 +1191,25 @@ def test_get_runtime_chains_parent_trace_context_for_workflow_tools( parent_external_trace_id="caller-trace", ), ) + assert tool_runtime.raw.parent_trace_context == chained assert workflow_runtime.runtime.runtime_parameters == {} + with ( + patch.object(ToolEngine, "generic_invoke", return_value=iter([])), + patch.object( + ToolFileMessageTransformer, + "transform_tool_invoke_messages", + side_effect=lambda *, messages, **_: messages, + ), + ): + list( + runtime.invoke( + tool_runtime=tool_runtime, tool_parameters={}, workflow_call_depth=0, provider_name="provider" + ) + ) + + workflow_runtime.set_parent_trace_context.assert_called_once_with(chained) + @staticmethod def test_get_runtime_passes_trace_session_and_caller_trace_to_workflow_tools( runtime: DifyToolNodeRuntime, From 59f27bc23a6f65425973dc5edfd8c1bd9841bf96 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Thu, 29 Oct 2026 05:31:41 +0800 Subject: [PATCH 06/17] revert(tracing): leave nested workflow trace placement for a follow-up This PR keeps the trace pipeline refactor as its goal. Revert the nested workflow placement (dfbccb85ba) and the test that covered its runtime hand-off (34286bc6ef). A nested workflow run keeps its own trace, linked to its outer run by attributes, as before. --- api/core/app/apps/advanced_chat/app_runner.py | 6 - api/core/app/apps/workflow/app_runner.py | 6 - api/core/app/apps/workflow_app_runner.py | 15 -- api/core/app/entities/app_invoke_entities.py | 9 - api/core/helper/trace_id_helper.py | 6 - api/core/tools/workflow_as_tool/tool.py | 12 +- api/core/workflow/node_runtime.py | 8 +- api/domain/tracing/nested_workflow_trace.py | 73 ------- api/domain/tracing/normalization.py | 6 +- api/domain/tracing/trace_run_agg.py | 2 - .../application_services/tracing.py | 8 +- api/extensions/tracing.py | 16 +- .../src/dify_trace_langfuse/provider.py | 5 +- .../src/dify_trace_langfuse/test_provider.py | 25 +-- .../src/dify_trace_langsmith/provider.py | 11 +- .../src/dify_trace_langsmith/test_provider.py | 20 -- .../src/dify_trace_mlflow/_spans.py | 1 - .../src/dify_trace_mlflow/test_provider.py | 17 -- .../src/dify_trace_opik/provider.py | 9 +- .../src/dify_trace_opik/test_provider.py | 16 -- .../trace-sdk/src/dify_trace_sdk/events.py | 20 +- .../trace-sdk/src/dify_trace_sdk/model.py | 12 +- .../trace-sdk/src/dify_trace_sdk/testing.py | 48 ----- .../src/dify_trace_sdk/test_emitter.py | 4 +- .../src/dify_trace_sdk/test_model.py | 28 --- .../src/dify_trace_sdk/test_otlp.py | 14 +- .../src/dify_trace_sdk/test_testing.py | 9 +- .../src/dify_trace_weave/provider.py | 2 +- .../src/dify_trace_weave/test_provider.py | 17 -- api/repositories/trace_record_repository.py | 45 +--- api/services/tracing/canonical.py | 14 +- api/services/tracing/delivery.py | 135 ++---------- api/services/tracing/records.py | 3 - api/services/tracing/recovery.py | 3 +- api/services/tracing/scopes.py | 4 +- api/services/tracing/service.py | 37 ++-- .../core/agent/test_cot_agent_runner.py | 4 +- .../core/agent/test_fc_agent_runner.py | 4 +- .../apps/advanced_chat/test_app_generator.py | 8 +- .../test_generate_task_pipeline.py | 7 +- .../app/apps/agent_app/test_app_generator.py | 8 +- .../app/apps/agent_chat/test_app_generator.py | 4 +- .../core/app/apps/chat/test_app_generator.py | 4 +- .../app/apps/completion/test_app_generator.py | 4 +- .../core/app/apps/test_base_app_runner.py | 4 +- .../core/app/apps/test_workflow_app_runner.py | 14 +- .../app/apps/workflow/test_app_generator.py | 16 +- .../app/entities/test_app_invoke_entities.py | 4 +- ...st_easy_ui_based_generate_task_pipeline.py | 8 +- .../core/llm_generator/test_llm_generator.py | 3 +- .../core/moderation/test_input_moderation.py | 4 +- .../rag/retrieval/test_dataset_retrieval.py | 4 +- .../core/telemetry/test___init__.py | 4 +- .../unit_tests/core/telemetry/test_gateway.py | 4 +- .../unit_tests/core/tools/test_tool_engine.py | 4 +- .../core/tools/workflow_as_tool/test_tool.py | 38 +--- .../core/workflow/test_node_runtime.py | 57 +----- .../tracing/test_nested_workflow_trace.py | 108 ---------- .../domain/tracing/test_normalization.py | 2 - .../domain/tracing/test_trace_run_agg.py | 1 - .../enterprise/telemetry/test_draft_trace.py | 4 +- .../telemetry/test_enterprise_trace.py | 1 - .../unit_tests/extensions/test_tracing.py | 35 +--- .../test_application_trace_repository.py | 33 ++- .../test_trace_record_repository.py | 143 +------------ .../services/test_workflow_persistence.py | 15 +- .../services/tracing/test_adapters.py | 5 +- .../services/tracing/test_canonical.py | 5 +- .../services/tracing/test_delivery.py | 193 +----------------- .../services/tracing/test_recovery.py | 8 +- .../services/tracing/test_runtime_bridge.py | 126 +++--------- .../services/tracing/test_scopes.py | 4 +- .../services/tracing/test_service.py | 5 +- .../tasks/test_application_trace_task.py | 4 +- .../tasks/test_trace_delivery_task.py | 26 +-- 75 files changed, 287 insertions(+), 1304 deletions(-) delete mode 100644 api/domain/tracing/nested_workflow_trace.py delete mode 100644 api/tests/unit_tests/domain/tracing/test_nested_workflow_trace.py diff --git a/api/core/app/apps/advanced_chat/app_runner.py b/api/core/app/apps/advanced_chat/app_runner.py index 5897360bb2fb2d..12687c89ee037b 100644 --- a/api/core/app/apps/advanced_chat/app_runner.py +++ b/api/core/app/apps/advanced_chat/app_runner.py @@ -132,8 +132,6 @@ def prepare(self) -> PreparedWorkflowRun | None: user_from=user_from, app_type=get_credit_usage_app_type(app_config.app_mode), trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), - external_trace_id=self.application_generate_entity.extras.get("external_trace_id"), - parent_trace_context=self.application_generate_entity.extras.get("parent_trace_context"), call_depth=self.application_generate_entity.call_depth, ) elif self.application_generate_entity.single_iteration_run or self.application_generate_entity.single_loop_run: @@ -145,8 +143,6 @@ def prepare(self) -> PreparedWorkflowRun | None: user_id=self.application_generate_entity.user_id, app_type=get_credit_usage_app_type(app_config.app_mode), trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), - external_trace_id=self.application_generate_entity.extras.get("external_trace_id"), - parent_trace_context=self.application_generate_entity.extras.get("parent_trace_context"), ) else: inputs = self.application_generate_entity.inputs @@ -220,8 +216,6 @@ def prepare(self) -> PreparedWorkflowRun | None: root_node_id=root_node_id, app_type=get_credit_usage_app_type(app_config.app_mode), trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), - external_trace_id=self.application_generate_entity.extras.get("external_trace_id"), - parent_trace_context=self.application_generate_entity.extras.get("parent_trace_context"), call_depth=self.application_generate_entity.call_depth, ) diff --git a/api/core/app/apps/workflow/app_runner.py b/api/core/app/apps/workflow/app_runner.py index 5171c6df4adb8d..b2390ea19b156c 100644 --- a/api/core/app/apps/workflow/app_runner.py +++ b/api/core/app/apps/workflow/app_runner.py @@ -97,8 +97,6 @@ def prepare(self) -> PreparedWorkflowRun | None: root_node_id=self._root_node_id, app_type=get_credit_usage_app_type(app_config.app_mode), trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), - external_trace_id=self.application_generate_entity.extras.get("external_trace_id"), - parent_trace_context=self.application_generate_entity.extras.get("parent_trace_context"), call_depth=self.application_generate_entity.call_depth, ) elif self.application_generate_entity.single_iteration_run or self.application_generate_entity.single_loop_run: @@ -109,8 +107,6 @@ def prepare(self) -> PreparedWorkflowRun | None: user_id=self.application_generate_entity.user_id, app_type=get_credit_usage_app_type(app_config.app_mode), trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), - external_trace_id=self.application_generate_entity.extras.get("external_trace_id"), - parent_trace_context=self.application_generate_entity.extras.get("parent_trace_context"), ) else: inputs = self.application_generate_entity.inputs @@ -157,8 +153,6 @@ def prepare(self) -> PreparedWorkflowRun | None: root_node_id=root_node_id, app_type=get_credit_usage_app_type(app_config.app_mode), trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), - external_trace_id=self.application_generate_entity.extras.get("external_trace_id"), - parent_trace_context=self.application_generate_entity.extras.get("parent_trace_context"), call_depth=self.application_generate_entity.call_depth, ) diff --git a/api/core/app/apps/workflow_app_runner.py b/api/core/app/apps/workflow_app_runner.py index 7fa0cef0dd683e..3df06d4b0df484 100644 --- a/api/core/app/apps/workflow_app_runner.py +++ b/api/core/app/apps/workflow_app_runner.py @@ -44,7 +44,6 @@ ) from core.app.layers.pause_state_persist_layer import PauseStateLayerConfig from core.credit_usage import CreditUsageAppType -from core.helper.trace_id_helper import ParentTraceContext from core.rag.entities import RetrievalSourceMetadata from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from core.repositories.human_input_repository import HumanInputFormRecord @@ -240,8 +239,6 @@ def _init_graph( root_node_id: str | None = None, app_type: CreditUsageAppType | None = None, trace_session_id: str | None = None, - external_trace_id: str | None = None, - parent_trace_context: ParentTraceContext | None = None, ) -> Graph: """ Init graph @@ -257,8 +254,6 @@ def _init_graph( invoke_from=invoke_from, app_type=app_type, trace_session_id=trace_session_id, - external_trace_id=external_trace_id, - parent_trace_context=parent_trace_context, ) graph_init_context = DifyGraphInitContext( execution_driver=self._execution_driver, @@ -295,8 +290,6 @@ def _prepare_single_node_execution( user_id: str, app_type: CreditUsageAppType | None = None, trace_session_id: str | None = None, - external_trace_id: str | None = None, - parent_trace_context: ParentTraceContext | None = None, ) -> tuple[Graph, VariablePool, RuntimeState]: """ Prepare graph, variable pool, and runtime state for single node execution @@ -335,8 +328,6 @@ def _prepare_single_node_execution( user_id=user_id, app_type=app_type, trace_session_id=trace_session_id, - external_trace_id=external_trace_id, - parent_trace_context=parent_trace_context, ) elif single_loop_run: graph, variable_pool = self._get_graph_and_variable_pool_for_single_node_run( @@ -348,8 +339,6 @@ def _prepare_single_node_execution( user_id=user_id, app_type=app_type, trace_session_id=trace_session_id, - external_trace_id=external_trace_id, - parent_trace_context=parent_trace_context, ) else: raise ValueError("Neither single_iteration_run nor single_loop_run is specified") @@ -369,8 +358,6 @@ def _get_graph_and_variable_pool_for_single_node_run( user_id: str = "", app_type: CreditUsageAppType | None = None, trace_session_id: str | None = None, - external_trace_id: str | None = None, - parent_trace_context: ParentTraceContext | None = None, ) -> tuple[Graph, VariablePool]: """ Get graph and variable pool for single node execution (iteration or loop). @@ -442,8 +429,6 @@ def _get_graph_and_variable_pool_for_single_node_run( invoke_from=InvokeFrom.DEBUGGER, app_type=app_type, trace_session_id=trace_session_id, - external_trace_id=external_trace_id, - parent_trace_context=parent_trace_context, ) graph_init_context = DifyGraphInitContext( execution_driver=self._execution_driver, diff --git a/api/core/app/entities/app_invoke_entities.py b/api/core/app/entities/app_invoke_entities.py index 98cb29ff9c7820..c2214eb088aabf 100644 --- a/api/core/app/entities/app_invoke_entities.py +++ b/api/core/app/entities/app_invoke_entities.py @@ -17,7 +17,6 @@ normalize_credit_usage_created_by, ) from core.entities.provider_configuration import ProviderModelBundle -from core.helper.trace_id_helper import ParentTraceContext from graphon.file import File, FileUploadConfig from graphon.model_runtime.entities.model_entities import AIModelEntity from models.model import AppMode @@ -103,10 +102,6 @@ class DifyRunContext(BaseModel): app_type: CreditUsageAppType | None = None created_by: CreditUsageCreatedBy | None = None trace_session_id: str | None = None - external_trace_id: str | None = None - """The trace ID the run's caller supplied.""" - parent_trace_context: ParentTraceContext | None = None - """The workflow tool node that started this run; ``None`` unless the run is nested.""" workflow_tool_invocation_id: str | None = None @field_validator("created_by", mode="before") @@ -134,8 +129,6 @@ def build_dify_run_context( app_type: CreditUsageAppTypeInput = None, created_by: CreditUsageCreatedByInput = None, trace_session_id: str | None = None, - external_trace_id: str | None = None, - parent_trace_context: ParentTraceContext | None = None, extra_context: Mapping[str, Any] | None = None, ) -> dict[str, Any]: """ @@ -154,8 +147,6 @@ def build_dify_run_context( app_type=normalize_credit_usage_app_type(app_type) if app_type is not None else None, created_by=normalize_credit_usage_created_by(created_by) if created_by is not None else None, trace_session_id=trace_session_id, - external_trace_id=external_trace_id, - parent_trace_context=parent_trace_context, ) return run_context diff --git a/api/core/helper/trace_id_helper.py b/api/core/helper/trace_id_helper.py index 41e4686909a058..e1ebd45e07478d 100644 --- a/api/core/helper/trace_id_helper.py +++ b/api/core/helper/trace_id_helper.py @@ -12,12 +12,6 @@ class ParentTraceContext(BaseModel): parent_workflow_run_id: StrictStr parent_node_execution_id: StrictStr | None = None - parent_app_id: StrictStr | None = None - """The app of the outer run; ``None`` in contexts saved before it was recorded.""" - parent_external_trace_id: StrictStr | None = None - """The trace ID the outer run's caller supplied; only an outermost run has one.""" - parent_trace_context: "ParentTraceContext | None" = None - """The outer run's own parent context; ``None`` when the outer run is not nested.""" model_config = ConfigDict(extra="forbid") diff --git a/api/core/tools/workflow_as_tool/tool.py b/api/core/tools/workflow_as_tool/tool.py index a9bf3ffb7a2bca..9598b4ebb6494a 100644 --- a/api/core/tools/workflow_as_tool/tool.py +++ b/api/core/tools/workflow_as_tool/tool.py @@ -255,9 +255,17 @@ def fork_tool_runtime(self, runtime: ToolRuntime) -> WorkflowTool: forked._trace_session_id = self._trace_session_id return forked - def set_parent_trace_context(self, context: ParentTraceContext) -> None: + def set_parent_trace_context( + self, + *, + parent_workflow_run_id: str, + parent_node_execution_id: str, + ) -> None: """Attach outer workflow trace context without exposing it as tool input.""" - self._parent_trace_context = context + self._parent_trace_context = ParentTraceContext( + parent_workflow_run_id=parent_workflow_run_id, + parent_node_execution_id=parent_node_execution_id, + ) def clear_parent_trace_context(self) -> None: """Remove parent trace context before invoking this tool outside a nested workflow.""" diff --git a/api/core/workflow/node_runtime.py b/api/core/workflow/node_runtime.py index 79dbebea35aceb..cd08dab797c9d4 100644 --- a/api/core/workflow/node_runtime.py +++ b/api/core/workflow/node_runtime.py @@ -610,9 +610,6 @@ def get_runtime( parent_trace_context = ParentTraceContext( parent_workflow_run_id=outer_workflow_run_id, parent_node_execution_id=node_execution_id, - parent_app_id=self._run_context.app_id, - parent_external_trace_id=self._run_context.external_trace_id, - parent_trace_context=self._run_context.parent_trace_context, ) if isinstance(self._run_context.trace_session_id, str) and self._run_context.trace_session_id: trace_session_id = self._run_context.trace_session_id @@ -650,7 +647,10 @@ def invoke( tool = runtime_binding.tool callback = DifyWorkflowCallbackHandler() if runtime_binding.parent_trace_context and hasattr(tool, "set_parent_trace_context"): - tool.set_parent_trace_context(runtime_binding.parent_trace_context) + tool.set_parent_trace_context( + parent_workflow_run_id=runtime_binding.parent_trace_context.parent_workflow_run_id, + parent_node_execution_id=runtime_binding.parent_trace_context.parent_node_execution_id, + ) elif hasattr(tool, "clear_parent_trace_context"): tool.clear_parent_trace_context() if runtime_binding.trace_session_id and hasattr(tool, "set_trace_session_id"): diff --git a/api/domain/tracing/nested_workflow_trace.py b/api/domain/tracing/nested_workflow_trace.py deleted file mode 100644 index b7788cab66b061..00000000000000 --- a/api/domain/tracing/nested_workflow_trace.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Place the trace of a nested workflow run in the trace of the run that started it. - -A nested run joins the trace of the run whose tool node started it when both -export to the same destination, and its root hangs from that tool node. That -outer run joined the trace above it by the same rule, so the joined trace is the -outermost one reachable through runs sharing the destination. A run whose spans -are collected for Agent tracing is always the root of its own trace. Every -identifier is derived from committed records, so the units need no coordination. -""" - -from collections.abc import Iterable -from dataclasses import dataclass - -from dify_trace_sdk.model import SpanRef, Trace - - -@dataclass(frozen=True, slots=True) -class OuterTrace: - """The trace an outer run's spans go to when that run is the root of its trace.""" - - trace_id: str - external_trace_id: str | None - """The caller-supplied trace ID the trace's root unit carries.""" - root: SpanRef - same_destination: bool - """Whether the nested run exports to the destination this trace goes to.""" - collected: bool - """Whether the run's spans are collected for Agent tracing, which keeps it the root of its trace.""" - - -@dataclass(frozen=True, slots=True) -class TracePlacement: - """Where the trace of a nested run hangs in an outer trace.""" - - trace_id: str - external_trace_id: str | None - parent: SpanRef - trace_root: SpanRef - - def place(self, trace: Trace) -> Trace: - """``trace`` as a unit of the outer trace, its root hanging from ``parent``.""" - root, *spans = trace.spans - return Trace( - trace_id=self.trace_id, - external_trace_id=self.external_trace_id, - session_id=trace.session_id, - user_id=trace.user_id, - parent=self.parent, - trace_root=self.trace_root, - spans=(root.model_copy(update={"parent_span_id": self.parent.span_id}), *spans), - ) - - -def nested_trace_placement(outer_traces: Iterable[OuterTrace], *, tool_node: SpanRef) -> TracePlacement | None: - """Where the trace of a nested run hangs; ``None`` when it starts a trace of its own. - - ``outer_traces`` follows the runs that started the nested run, innermost - first, and ends at the first run whose trace is not known yet. It is consumed - only as far as the decision needs. ``tool_node`` is the span of the tool node - that started the nested run. - """ - joined: OuterTrace | None = None - for outer in outer_traces: - if not outer.same_destination: - break - joined = outer - if outer.collected: - break - if joined is None: - return None - return TracePlacement( - trace_id=joined.trace_id, external_trace_id=joined.external_trace_id, parent=tool_node, trace_root=joined.root - ) diff --git a/api/domain/tracing/normalization.py b/api/domain/tracing/normalization.py index e7a2c337fb7b53..aba79110d456ff 100644 --- a/api/domain/tracing/normalization.py +++ b/api/domain/tracing/normalization.py @@ -278,6 +278,10 @@ def normalize_agent_span(data: Mapping[str, Any], *, capture_content: bool) -> S "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", } ) @@ -313,7 +317,6 @@ def remap_outer_trace( session_id=trace.session_id, user_id=trace.user_id, parent=None, - trace_root=None, spans=tuple(spans), ) @@ -360,7 +363,6 @@ def legacy_outer_trace(data: Mapping[str, Any]) -> Trace: session_id=data.get("session_id") or None, user_id=None, parent=None, - trace_root=None, spans=tuple(spans), ) diff --git a/api/domain/tracing/trace_run_agg.py b/api/domain/tracing/trace_run_agg.py index 0519f83f1acc72..f836154cc726e4 100644 --- a/api/domain/tracing/trace_run_agg.py +++ b/api/domain/tracing/trace_run_agg.py @@ -422,7 +422,6 @@ def settle_preview(self, node: SourceNodeExecution, *, now: datetime) -> None: session_id=self.session_id, user_id=None, parent=None, - trace_root=None, spans=(span,), ) self._keep_terminal(trace, now=now) @@ -639,6 +638,5 @@ def exported_trace(self, batches: Sequence[TraceBatch]) -> Trace: session_id=self.session_id, user_id=outer.user_id, parent=None, - trace_root=None, spans=order_spans(outer.root.span_id, spans), ) diff --git a/api/extensions/application_services/tracing.py b/api/extensions/application_services/tracing.py index aaf0000d63e1df..26e8b3d61ecc9a 100644 --- a/api/extensions/application_services/tracing.py +++ b/api/extensions/application_services/tracing.py @@ -46,9 +46,6 @@ def build_tracing_services(*, database_client: sessionmaker[Session]) -> Tracing builder = CanonicalTraceBuilder() configs = SQLAlchemyAppTracingConfigRepository(session_factory=database_client) enterprise_telemetry = is_enterprise_telemetry_enabled() - records = TraceRecordRepository( - database_client, files_url=dify_config.FILES_URL, enterprise_details=enterprise_telemetry - ) agent_settings = AgentTraceSettings( enabled=dify_config.OPS_AGENT_TRACE_ENABLED, max_spans=dify_config.OPS_AGENT_TRACE_MAX_SPANS, @@ -77,8 +74,9 @@ def build_tracing_services(*, database_client: sessionmaker[Session]) -> Tracing emitter=CeleryTraceEmitter(), policies=configs, enterprise_enabled=is_enterprise_telemetry_enabled ), delivery=TraceDeliveryService( - records=records, - outer_runs=records, + records=TraceRecordRepository( + database_client, files_url=dify_config.FILES_URL, enterprise_details=enterprise_telemetry + ), destinations=configs, providers=providers, codec=codec, diff --git a/api/extensions/tracing.py b/api/extensions/tracing.py index a1512d499e0bb6..93e60916e43063 100644 --- a/api/extensions/tracing.py +++ b/api/extensions/tracing.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Literal, Protocol, cast from dify_trace_sdk.emitter import TraceScope -from dify_trace_sdk.events import TraceCorrelation, TraceParent +from dify_trace_sdk.events import TraceCorrelation from flask import Flask, current_app, has_app_context from core.helper.trace_id_helper import ParentTraceContext @@ -88,17 +88,6 @@ def _scopes() -> ExecutionTraceScopes | None: return cast(ExecutionTraceScopes | None, current_app.extensions.get(_SCOPES_KEY)) -def _trace_parent(context: ParentTraceContext) -> TraceParent: - outer = context.parent_trace_context - return TraceParent( - workflow_run_id=context.parent_workflow_run_id, - app_id=context.parent_app_id, - node_execution_id=context.parent_node_execution_id, - external_trace_id=context.parent_external_trace_id, - parent=_trace_parent(outer) if outer is not None 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") @@ -111,7 +100,8 @@ def correlation_from_extras(extras: Mapping[str, Any]) -> TraceCorrelation: 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=_trace_parent(parent) if isinstance(parent, ParentTraceContext) 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, ) 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 index b275b49f13a01c..b587bba7c8d837 100644 --- a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/provider.py +++ b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/provider.py @@ -118,7 +118,7 @@ def _observation_id(trace: Trace, span_id: str) -> str: def _events(self, trace: Trace, trace_key: str) -> list[dict[str, object]]: events: list[dict[str, object]] = [] - if trace.trace_root is None: + if trace.parent is None: root = trace.root events.append( self._event( @@ -140,8 +140,7 @@ def _events(self, trace: Trace, trace_key: str) -> list[dict[str, object]]: trace_level_span_id = root.span_id observations: Sequence[Span] = trace.spans[1:] else: - # Spans under the trace root, which another unit sent, are top-level observations. - trace_level_span_id = trace.trace_root.span_id + trace_level_span_id = trace.parent.span_id observations = trace.spans for span in observations: parent = ( 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 index 0b6227958915a6..0c04bf5edb88cc 100644 --- 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 @@ -11,16 +11,7 @@ VerificationError, ) from dify_trace_sdk.ids import span_uuid -from dify_trace_sdk.testing import ( - MESSAGE_ID, - WORKFLOW_RUN_ID, - RecordingEgress, - message_trace, - nested_workflow_fragment, - response, - tool_fragment, - workflow_trace, -) +from dify_trace_sdk.testing import MESSAGE_ID, RecordingEgress, message_trace, response, tool_fragment, workflow_trace CONFIG = LangfuseConfig(public_key="pk-lf-1", secret_key="sk-lf-1", host="https://cloud.langfuse.com") @@ -109,20 +100,6 @@ def test_fragment_joins_the_message_trace_without_recreating_it() -> None: assert "parentObservationId" not in tool["body"] -def test_nested_workflow_hangs_from_the_tool_node_that_started_it() -> None: - provider, egress = _provider(response(207, {"successes": [{"id": "x", "status": 201}], "errors": []})) - fragment = nested_workflow_fragment() - assert fragment.parent is not None - - provider.send(CONFIG, fragment) - - run, node = _sent_events(egress) - assert [run["type"], node["type"]] == ["span-create", "span-create"] - assert run["body"]["traceId"] == node["body"]["traceId"] == WORKFLOW_RUN_ID - assert run["body"]["parentObservationId"] == str(span_uuid(fragment.trace_id, fragment.parent.span_id)) - assert node["body"]["parentObservationId"] == run["body"]["id"] - - 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"}) 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 index 16b792e71a97c7..aa48e92110fd80 100644 --- a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/provider.py +++ b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/provider.py @@ -137,7 +137,7 @@ def _trace_run_id(trace: Trace) -> str: 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.trace_root.span_id if trace.trace_root is not None else trace.root.span_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)) @@ -251,15 +251,12 @@ def send(self, config: LangSmithConfig, trace: Trace) -> SendResult: 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 or trace.trace_root is None: + 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. - # A deeper parent's dotted order names every run above it, so the unit hangs from the root. - root_order = _dotted_order(root_run_id, trace.trace_root.start_time, None) - for anchor in (trace.trace_root.span_id, trace.parent.span_id): - run_ids[anchor] = root_run_id - dotted_orders[anchor] = root_order + 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: 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 index c1fa8496cd7789..10ed1402d3b1f7 100644 --- 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 @@ -20,7 +20,6 @@ WORKFLOW_RUN_ID, RecordingEgress, message_trace, - nested_workflow_fragment, response, tool_fragment, workflow_trace, @@ -278,25 +277,6 @@ def test_fragment_attaches_under_the_message_run() -> None: ) -def test_nested_workflow_attaches_under_the_trace_root_run() -> None: - """LangSmith orders runs by their full ancestry, which a deeper parent's unit does not carry.""" - workflow_provider, workflow_egress = _provider(*_accepted(3)) - workflow_provider.send(VERIFIED, workflow_trace()) - root_order = _runs(workflow_egress)[0]["dotted_order"] - provider, egress = _provider(*_accepted(2)) - fragment = nested_workflow_fragment() - run_id = str(span_uuid(fragment.trace_id, fragment.root.span_id)) - - provider.send(VERIFIED, fragment) - - nested_run, node_run = _runs(egress) - assert nested_run["id"] == run_id - assert (nested_run["parent_run_id"], nested_run["trace_id"]) == (WORKFLOW_RUN_ID, WORKFLOW_RUN_ID) - assert nested_run["dotted_order"] == f"{root_order}.20261001T080003100000Z{run_id}" - assert node_run["parent_run_id"] == run_id - assert node_run["dotted_order"].startswith(f"{nested_run['dotted_order']}.") - - 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/") 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 index 1ab12d34a84dce..aee92de2426076 100644 --- a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/_spans.py +++ b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/_spans.py @@ -55,7 +55,6 @@ def standalone(trace: Trace) -> Trace: update={ "trace_id": trace_id_for(f"{trace.trace_id}:{root.span_id}"), "parent": None, - "trace_root": None, "spans": (detached, *trace.spans[1:]), } ) 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 index e1ea2132bbf565..9624c813690de2 100644 --- 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 @@ -28,7 +28,6 @@ RecordedRequest, RecordingEgress, message_trace, - nested_workflow_fragment, response, tool_fragment, workflow_trace, @@ -350,22 +349,6 @@ def test_fragment_is_sent_as_its_own_trace_linked_to_its_parent() -> None: assert result.trace_id == f"tr-{own_trace_id}" assert own_trace_id != fragment.trace_id - @staticmethod - def test_nested_workflow_is_its_own_trace_naming_the_tool_node_it_hangs_from() -> None: - provider, egress = _mlflow_provider(ACK) - fragment = nested_workflow_fragment() - assert fragment.parent is not None - own_trace_id = trace_id_for(f"{fragment.trace_id}:{fragment.root.span_id}") - - provider.send(MLFLOW_CONFIG, fragment) - - run, node = _exported(egress) - assert (run.trace_id, node.trace_id) == (f"tr-{own_trace_id}", f"tr-{own_trace_id}") - assert run.parent_id is None - assert run.get_attribute("dify.parent_trace_id") == fragment.trace_id - assert run.get_attribute("dify.parent_span_id") == fragment.parent.span_id - assert node.parent_id == run.span_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"}) 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 index 4e1639629924f2..a1f6de4a5dd229 100644 --- a/api/providers/trace/trace-opik/src/dify_trace_opik/provider.py +++ b/api/providers/trace/trace-opik/src/dify_trace_opik/provider.py @@ -226,8 +226,8 @@ def send(self, config: OpikConfig, trace: Trace) -> SendResult: settings = _settings(config) except VerificationError as error: raise ExportRejectedError(f"Opik configuration is invalid ({error.code})") from None - # A fragment names the root of its trace, so both units derive the same trace ID. - trace_start = trace.trace_root.start_time if trace.trace_root is not None else trace.root.start_time + # 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) @@ -278,9 +278,6 @@ def _check(self, method: str, url: str, headers: Mapping[str, str], content: byt 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} - if trace.parent is not None and trace.parent != trace.trace_root: - # A parent below the trace root is a span another unit sent; its ID is derivable. - span_ids[trace.parent.span_id] = _opik_id(trace.trace_id, trace.parent.span_id, trace.parent.start_time) 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: @@ -315,7 +312,7 @@ def _span(trace: Trace, span: Span, trace_id: str, span_ids: Mapping[str, str], "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 that hangs from the trace root. + # 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"), 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 index a968a87f0de71c..00a9890e8f6022 100644 --- 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 @@ -24,7 +24,6 @@ RecordedRequest, RecordingEgress, message_trace, - nested_workflow_fragment, response, tool_fragment, workflow_trace, @@ -457,7 +456,6 @@ def test_large_units_send_spans_in_batches() -> None: session_id=None, user_id=None, parent=None, - trace_root=None, spans=(run, *nodes), ) provider, egress = _provider(CREATED, NO_CONTENT, NO_CONTENT) @@ -516,20 +514,6 @@ def test_span_failures_after_the_trace_was_created_are_outcome_unknown(outcome: assert len(egress.requests) == 2 -def test_nested_workflow_hangs_from_the_tool_node_span_of_its_outer_trace() -> None: - provider, egress = _provider(NO_CONTENT) - - assert provider.send(CONFIG, nested_workflow_fragment()).trace_id == WORKFLOW_TRACE_ID - - (create_spans,) = egress.requests - assert create_spans.url == f"{API_URL}v1/private/spans/batch" - body = create_spans.json() - assert isinstance(body, dict) - run, node = body["spans"] - assert (run["trace_id"], run["parent_span_id"]) == (WORKFLOW_TRACE_ID, WORKFLOW_TOOL_SPAN_ID) - assert (node["trace_id"], node["parent_span_id"]) == (WORKFLOW_TRACE_ID, run["id"]) - - @pytest.mark.parametrize( ("outcome", "expected"), [(response(503, {}), ExportRetryableError), (response(400, {}), ExportRejectedError)], 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 index 4a55503b75f94b..3f1b8dd3ac3714 100644 --- a/api/providers/trace/trace-sdk/src/dify_trace_sdk/events.py +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/events.py @@ -20,29 +20,13 @@ def json_safe(value: object) -> JsonValue: return json.loads(json.dumps(value, ensure_ascii=False, default=str), parse_constant=str) -class TraceParent(BaseModel): - """The workflow run, and its tool node execution, that started a nested workflow run.""" - - workflow_run_id: str - app_id: str | None - """``None`` when the context that started the nested run did not record it.""" - node_execution_id: str | None - """``None`` when the starting node execution is unknown.""" - external_trace_id: str | None - """The trace ID that run's caller supplied; only an outermost run has one.""" - parent: "TraceParent | None" - """The run that started that run in turn; ``None`` when that run is not nested.""" - - model_config = _EVENT_CONFIG - - class TraceCorrelation(BaseModel): """Caller-supplied correlation, propagated unchanged to the trace.""" external_trace_id: str | None session_id: str | None - parent: TraceParent | None - """The run that started this nested workflow run; ``None`` outside nested workflows.""" + parent_workflow_run_id: str | None + parent_node_execution_id: str | None model_config = _EVENT_CONFIG 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 index adfe42c630725a..3203233f084595 100644 --- a/api/providers/trace/trace-sdk/src/dify_trace_sdk/model.py +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/model.py @@ -153,10 +153,8 @@ 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: a tool call under its message, or a nested workflow run under - the tool node that started it. The first span of this unit points at it. - ``trace_root`` is the root span of that trace, the same span as ``parent`` - when the unit hangs directly from it. + 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 @@ -165,8 +163,6 @@ class Trace(BaseModel): session_id: str | None user_id: str | None parent: SpanRef | None - trace_root: SpanRef | None - """Set exactly when ``parent`` is set.""" spans: tuple[Span, ...] = Field(min_length=1) model_config = ConfigDict(extra="forbid", frozen=True) @@ -191,10 +187,6 @@ def _parent_first(self) -> Self: seen.add(span.span_id) if expected_root_parent in seen: raise ValueError("the external parent cannot be a span of this unit") - if (self.parent is None) != (self.trace_root is None): - raise ValueError("parent and trace_root must be set together") - if self.trace_root is not None and self.trace_root.span_id in seen: - raise ValueError("the trace root cannot be a span of this unit") return self @property 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 index 60fde5fb60f468..846d38716062a4 100644 --- a/api/providers/trace/trace-sdk/src/dify_trace_sdk/testing.py +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/testing.py @@ -15,7 +15,6 @@ MESSAGE_ID = "7b3c6c0e-38a1-4b52-9a7e-3ad8f3a1c001" WORKFLOW_RUN_ID = "7b3c6c0e-38a1-4b52-9a7e-3ad8f3a1c002" -NESTED_WORKFLOW_RUN_ID = "7b3c6c0e-38a1-4b52-9a7e-3ad8f3a1c003" STARTED_AT = datetime(2026, 10, 1, 8, 0, 0, tzinfo=UTC) @@ -90,7 +89,6 @@ def message_trace() -> Trace: session_id="conversation-1", user_id="end-user-1", parent=None, - trace_root=None, spans=( Span( span_id=message_span_id, @@ -135,7 +133,6 @@ def workflow_trace() -> Trace: session_id=WORKFLOW_RUN_ID, user_id="account-1", parent=None, - trace_root=None, spans=( Span( span_id=run_span_id, @@ -190,7 +187,6 @@ def tool_fragment() -> Trace: session_id="conversation-1", user_id="end-user-1", parent=SpanRef(span_id=span_id_for(MESSAGE_ID), start_time=STARTED_AT), - trace_root=SpanRef(span_id=span_id_for(MESSAGE_ID), start_time=STARTED_AT), spans=( Span( span_id=span_id_for("tool-event-1"), @@ -206,47 +202,3 @@ def tool_fragment() -> Trace: ), ), ) - - -def nested_workflow_fragment() -> Trace: - """A nested workflow run sent separately and attached under the tool node of ``workflow_trace``.""" - tool_span_id = span_id_for("node-execution-tool") - run_span_id = span_id_for(NESTED_WORKFLOW_RUN_ID) - return Trace( - trace_id=trace_id_for(WORKFLOW_RUN_ID), - external_trace_id=None, - session_id=WORKFLOW_RUN_ID, - user_id="account-1", - parent=SpanRef(span_id=tool_span_id, start_time=STARTED_AT + timedelta(seconds=3)), - trace_root=SpanRef(span_id=span_id_for(WORKFLOW_RUN_ID), start_time=STARTED_AT), - spans=( - Span( - span_id=run_span_id, - parent_span_id=tool_span_id, - name="workflow", - kind=SpanKind.CHAIN, - start_time=STARTED_AT + timedelta(seconds=3, milliseconds=100), - end_time=STARTED_AT + timedelta(seconds=4), - status=SpanStatus.OK, - inputs={"query": "tracing"}, - outputs={"answer": "Tracing records execution."}, - attributes={ - "workflow_run_id": NESTED_WORKFLOW_RUN_ID, - "parent_workflow_run_id": WORKFLOW_RUN_ID, - "parent_node_execution_id": "node-execution-tool", - }, - ), - Span( - span_id=span_id_for("nested-node-execution-code"), - parent_span_id=run_span_id, - name="code_Lookup", - kind=SpanKind.CHAIN, - start_time=STARTED_AT + timedelta(seconds=3, milliseconds=200), - end_time=STARTED_AT + timedelta(seconds=3, milliseconds=900), - status=SpanStatus.OK, - inputs={"query": "tracing"}, - outputs={"answer": "Tracing records execution."}, - attributes={"node_execution_id": "nested-node-execution-code", "node_type": "code"}, - ), - ), - ) 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 index 42d5428173b154..67b5b7b56afc2d 100644 --- 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 @@ -4,7 +4,9 @@ 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=None) +CORRELATION = TraceCorrelation( + external_trace_id="caller-1", session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None +) class _CollectingEmitter: 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 index ad7f57dfe31253..a66539abd6a8d6 100644 --- 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 @@ -59,7 +59,6 @@ def test_trace_requires_parent_first_order() -> None: session_id=None, user_id=None, parent=None, - trace_root=None, spans=(_span(root, None), _span(span_id_for("orphan"), child)), ) with pytest.raises(ValidationError, match="root of this unit"): @@ -69,7 +68,6 @@ def test_trace_requires_parent_first_order() -> None: session_id=None, user_id=None, parent=None, - trace_root=None, spans=(_span(child, root),), ) @@ -83,36 +81,10 @@ def test_fragment_root_must_point_at_its_external_parent() -> None: session_id=None, user_id=None, parent=parent, - trace_root=parent, spans=(_span(span_id_for("tool"), None),), ) -def test_fragment_names_its_parent_and_the_root_of_its_trace_together() -> None: - trace_root = SpanRef(span_id=span_id_for("workflow"), start_time=STARTED_AT) - tool_node = SpanRef(span_id=span_id_for("tool-node"), start_time=STARTED_AT) - nested = _span(span_id_for("nested-workflow"), tool_node.span_id) - - def fragment(parent: SpanRef | None, root: SpanRef | None, spans: tuple[Span, ...]) -> Trace: - return Trace( - trace_id=trace_id_for("t"), - external_trace_id=None, - session_id=None, - user_id=None, - parent=parent, - trace_root=root, - spans=spans, - ) - - assert fragment(tool_node, trace_root, (nested,)).trace_root == trace_root - with pytest.raises(ValidationError, match="set together"): - fragment(tool_node, None, (nested,)) - with pytest.raises(ValidationError, match="set together"): - fragment(None, trace_root, (_span(span_id_for("root"), None),)) - with pytest.raises(ValidationError, match="trace root cannot be a span of this unit"): - fragment(tool_node, SpanRef(span_id=nested.span_id, start_time=STARTED_AT), (nested,)) - - def test_cost_must_be_finite() -> None: with pytest.raises(ValidationError): Cost(amount=float("inf"), currency="USD") 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 index f54c457d55aa85..c72e06a793e7d8 100644 --- 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 @@ -10,7 +10,7 @@ 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, nested_workflow_fragment, response, tool_fragment +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 @@ -26,18 +26,6 @@ def test_otlp_encoding_keeps_ids_and_external_parent() -> None: assert [attribute.key for attribute in span.attributes] == ["a"] -def test_otlp_encoding_hangs_a_nested_unit_from_the_span_that_started_it() -> None: - fragment = nested_workflow_fragment() - assert fragment.parent is not None - - request = encode_trace(fragment, resource={"service.name": "dify"}, attributes=lambda span: {}) - - run, node = request.resource_spans[0].scope_spans[0].spans - assert run.trace_id.hex() == node.trace_id.hex() == fragment.trace_id - assert run.parent_span_id.hex() == fragment.parent.span_id - assert node.parent_span_id.hex() == fragment.root.span_id - - def test_otlp_ack_partial_success_is_outcome_unknown() -> None: partial = ExportTraceServiceResponse() partial.partial_success.rejected_spans = 1 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 index d7b86dfc91e534..7a97ac9c91c363 100644 --- 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 @@ -1,4 +1,4 @@ -from dify_trace_sdk.testing import message_trace, nested_workflow_fragment, tool_fragment, workflow_trace +from dify_trace_sdk.testing import message_trace, tool_fragment, workflow_trace def test_sample_traces_are_valid_parent_first_trees() -> None: @@ -7,10 +7,3 @@ def test_sample_traces_are_valid_parent_first_trees() -> None: fragment = tool_fragment() assert fragment.parent is not None assert fragment.root.parent_span_id == fragment.parent.span_id - nested = nested_workflow_fragment() - assert nested.parent is not None - assert nested.trace_root is not None - assert nested.root.parent_span_id == nested.parent.span_id - assert nested.trace_id == workflow_trace().trace_id - assert nested.trace_root.span_id == workflow_trace().root.span_id - assert nested.parent.span_id in {span.span_id for span in workflow_trace().spans} 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 index d962fd1de58c18..9554d3bf201a20 100644 --- a/api/providers/trace/trace-weave/src/dify_trace_weave/provider.py +++ b/api/providers/trace/trace-weave/src/dify_trace_weave/provider.py @@ -281,7 +281,7 @@ def send(self, config: WeaveConfig, trace: Trace) -> SendResult: 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.trace_root.span_id if trace.trace_root is not None else trace.root.span_id + 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)}", 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 index a065a4e592fa2c..6253d07e071a19 100644 --- 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 @@ -22,7 +22,6 @@ RecordedRequest, RecordingEgress, message_trace, - nested_workflow_fragment, response, tool_fragment, workflow_trace, @@ -478,22 +477,6 @@ def test_fragment_joins_the_message_call_without_naming_a_turn() -> None: ) -def test_nested_workflow_hangs_from_the_tool_node_call_that_started_it() -> None: - fragment = nested_workflow_fragment() - assert fragment.parent is not None - assert fragment.trace_root is not None - provider, egress = _provider(*_acknowledgments(fragment)) - - result = provider.send(CONFIG, fragment) - - run_start, node_start = (_body(request)["start"] for request in egress.requests[:2]) - assert (run_start["trace_id"], node_start["trace_id"]) == (WORKFLOW_RUN_ID, WORKFLOW_RUN_ID) - assert run_start["parent_id"] == _call_id(fragment, fragment.parent.span_id) - assert node_start["parent_id"] == run_start["id"] - assert run_start["turn_id"] is None - assert result.trace_url == f"{CALLS_URL}/{_call_id(fragment, fragment.trace_root.span_id)}" - - 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)) diff --git a/api/repositories/trace_record_repository.py b/api/repositories/trace_record_repository.py index 9a19a994995e05..0b6bd57a1f3522 100644 --- a/api/repositories/trace_record_repository.py +++ b/api/repositories/trace_record_repository.py @@ -31,13 +31,13 @@ from graphon.enums import WorkflowExecutionStatus from models import Account, App, Tenant from models.dataset import Dataset -from models.model import AppMode, Conversation, EndUser, Message, MessageFile +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, WorkflowNodeExecutionTriggeredFrom, WorkflowRun from repositories.factory import DifyAPIRepositoryFactory from repositories.tracing_store_errors import tracing_store_errors -from services.tracing.delivery import OuterRunQuery, OuterRunRoot, TraceRecordSource +from services.tracing.delivery import TraceRecordSource from services.tracing.errors import SourceNotFoundError, SourceNotReadyError from services.tracing.records import ( BaseTraceInfo, @@ -123,7 +123,7 @@ def _streaming_metrics(message: Message) -> dict[str, Any]: } -class TraceRecordRepository(TraceRecordSource, OuterRunQuery): +class TraceRecordRepository(TraceRecordSource): def __init__(self, session_factory: sessionmaker[Session], *, files_url: str, enterprise_details: bool) -> None: self._sessions = session_factory self._files_url = files_url @@ -154,38 +154,6 @@ def assemble(self, event: TraceEvent) -> BaseTraceInfo: # Shared lookups - @override - def root(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> OuterRunRoot | None: - runs = DifyAPIRepositoryFactory.create_api_workflow_run_repository(self._sessions) - with tracing_store_errors(): - run = runs.get_workflow_run_by_id(tenant_id=tenant_id, app_id=app_id, run_id=workflow_run_id) - if run is None: - return None - with tracing_store_errors(), self._sessions() as session: - mode = session.scalar(select(App.mode).where(App.id == app_id, App.tenant_id == tenant_id)) - if mode is not AppMode.ADVANCED_CHAT: - return OuterRunRoot(root_id=run.id, started_at=run.created_at) - message = session.execute( - select(Message.id, Message.created_at).where( - Message.app_id == app_id, Message.workflow_run_id == run.id - ) - ).first() - if message is None: - return None - message_id, created_at = message - return OuterRunRoot(root_id=message_id, started_at=created_at) - - @override - def node_started_at( - self, *, tenant_id: str, app_id: str, workflow_run_id: str, node_execution_id: str - ) -> datetime | None: - executions = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(self._sessions) - with tracing_store_errors(): - execution = executions.get_execution_by_id(execution_id=node_execution_id, tenant_id=tenant_id) - if execution is None or execution.app_id != app_id or execution.workflow_run_id != workflow_run_id: - return None - return execution.created_at - 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: @@ -362,10 +330,10 @@ def _workflow(self, event: TraceEvent, payload: WorkflowRunCompleted) -> Workflo "workspace_name": workspace_name, } correlation = event.correlation - if correlation.parent is not None: + 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, + "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, @@ -392,7 +360,6 @@ def _workflow(self, event: TraceEvent, payload: WorkflowRunCompleted) -> Workflo start_time=run.created_at, end_time=run.finished_at, invoked_by=_user_of(metadata), - parent=correlation.parent, ) def _node_executions(self, run: WorkflowRun) -> Sequence[WorkflowNodeExecution]: diff --git a/api/services/tracing/canonical.py b/api/services/tracing/canonical.py index 83806c5ffe82be..505fb01a7a645b 100644 --- a/api/services/tracing/canonical.py +++ b/api/services/tracing/canonical.py @@ -43,9 +43,8 @@ _RETRY_SUMMARY_FIELDS = ("retry_index", "error", "elapsed_time", "created_at", "finished_at") -def root_trace_id(external_trace_id: str | None, root_id: str) -> str: - """The trace ID of a trace rooted at the message or workflow run ``root_id``.""" - return trace_id_for(external_trace_id or root_id) +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: @@ -180,12 +179,11 @@ def _message(self, record: MessageTraceInfo) -> Trace: status = _status(error, None) attributes = _attributes({**record.metadata, "trace_entity_type": "message"}) return Trace( - trace_id=root_trace_id(record.trace_id, message.id), + 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, - trace_root=None, spans=( Span( span_id=root_span_id, @@ -325,12 +323,11 @@ def _workflow(self, record: WorkflowTraceInfo) -> Trace: } ) return Trace( - trace_id=root_trace_id(record.trace_id, root_logical_id), + 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, - trace_root=None, spans=tuple(spans), ) @@ -419,7 +416,7 @@ def _auxiliary( parent = None parent_span_id = None else: - trace_id = root_trace_id(record.trace_id, message.id) + 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( @@ -428,7 +425,6 @@ def _auxiliary( session_id=_session_id(record, message.conversation_id if message is not None else None), user_id=_user_id(record), parent=parent, - trace_root=parent, spans=( Span( span_id=span_id, diff --git a/api/services/tracing/delivery.py b/api/services/tracing/delivery.py index d24bb75e8827dc..9e91e3e0a80d22 100644 --- a/api/services/tracing/delivery.py +++ b/api/services/tracing/delivery.py @@ -2,36 +2,31 @@ 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. A nested workflow run joins the trace of -the run that started it when both go to the same destination. 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. +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, Iterator, Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass -from datetime import datetime 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, TraceParent -from dify_trace_sdk.ids import span_id_for -from dify_trace_sdk.model import SpanRef, Trace +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 domain.tracing.nested_workflow_trace import OuterTrace, nested_trace_placement -from services.tracing.canonical import CanonicalTraceBuilder, root_trace_id +from services.tracing.canonical import CanonicalTraceBuilder from services.tracing.errors import ( ProviderConfigUnreadableError, TraceCollectorUnavailableError, TraceProviderUnavailableError, - TracingStoreUnavailableError, UnknownTraceProviderError, ) from services.tracing.providers import ProviderConfigCodec, TraceProviderCatalog @@ -72,46 +67,13 @@ def active_destination(self, *, tenant_id: str, app_id: str) -> StoredTraceDesti ... -@dataclass(frozen=True, slots=True) -class OuterRunRoot: - """The span a workflow run's trace starts with when that run is the root of its trace.""" - - root_id: str - """The chatflow message the run answers, otherwise the run itself.""" - started_at: datetime - - -class OuterRunQuery(Protocol): - """The runs that started nested workflow runs, read through their tenant and app.""" - - def root(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> OuterRunRoot | None: - """``None`` until the run, and the message a chatflow run answers, are persisted.""" - ... - - def node_started_at( - self, *, tenant_id: str, app_id: str, workflow_run_id: str, node_execution_id: str - ) -> datetime | None: - """``None`` until the node execution is persisted.""" - ... - - -@dataclass(frozen=True, slots=True) -class CollectedTrace: - """How the Agent trace of a workflow run is exported.""" - - trace_id: str | None - """``None`` until collection starts and fixes it.""" - destination: StoredTraceDestination | None - """The destination frozen when collection started; ``None`` when the run exports nowhere.""" - - 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 collected_trace(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> CollectedTrace | None: - """How the run's Agent trace is exported; ``None`` when its spans are not collected.""" + 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.""" ... @@ -141,10 +103,6 @@ class Destination: provider: TraceProvider[Any] config: ProviderConfig - def is_same(self, other: "Destination") -> bool: - """Whether both send to the same provider with the same configuration.""" - return self.provider.id == other.provider.id and self.config == other.config - def retry_delay_seconds(attempt: int) -> int: return min(600, 10 * 2 ** (attempt - 1)) @@ -155,7 +113,6 @@ def __init__( self, *, records: TraceRecordSource, - outer_runs: OuterRunQuery, destinations: TraceDestinationQuery, providers: TraceProviderCatalog, codec: ProviderConfigCodec, @@ -168,7 +125,6 @@ def __init__( ) -> None: """``enterprise_sink`` and ``collector`` are ``None`` when the deployment does not run that sink.""" self._records = records - self._outer_runs = outer_runs self._destinations = destinations self._providers = providers self._codec = codec @@ -198,6 +154,8 @@ def deliver(self, event: TraceEvent, *, enterprise_delivered: bool) -> None: 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: @@ -220,8 +178,6 @@ def deliver(self, event: TraceEvent, *, enterprise_delivered: bool) -> None: if trace is None: _event_counter.add(1, {"kind": kind, "outcome": "not_exported"}) return - if isinstance(record, WorkflowTraceInfo) and record.parent is not None: - trace = self._nested(event.tenant_id, record.parent, destination, trace) _event_counter.add(1, {"kind": kind, "outcome": "exported"}) self.export( ExportJob( @@ -276,71 +232,8 @@ def export(self, job: ExportJob, destination: Destination) -> None: 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) - return self._decoded(tenant_id, app_id, stored) if stored is not None else None - - def _nested(self, tenant_id: str, parent: TraceParent, destination: Destination, trace: Trace) -> Trace: - """Hang a nested run's trace from the tool node that started it when its outer trace goes to ``destination``. - - A nested run whose outer records cannot be read yet is sent on its own. - """ - if parent.app_id is None or parent.node_execution_id is None: - return trace - try: - started_at = self._outer_runs.node_started_at( - tenant_id=tenant_id, - app_id=parent.app_id, - workflow_run_id=parent.workflow_run_id, - node_execution_id=parent.node_execution_id, - ) - if started_at is None: - return trace - placement = nested_trace_placement( - self._outer_traces(tenant_id, parent, destination), - tool_node=SpanRef(span_id=span_id_for(parent.node_execution_id), start_time=started_at), - ) - except TracingStoreUnavailableError: - logger.warning("Outer trace unreadable; nested trace sent on its own: app_id=%s", parent.app_id) - return trace - return placement.place(trace) if placement is not None else trace - - def _outer_traces(self, tenant_id: str, parent: TraceParent, destination: Destination) -> Iterator[OuterTrace]: - """The traces of the runs that started a nested run, innermost first, while they are known.""" - link: TraceParent | None = parent - while link is not None and link.app_id is not None: - root = self._outer_runs.root(tenant_id=tenant_id, app_id=link.app_id, workflow_run_id=link.workflow_run_id) - if root is None: - return - collected = ( - self._collector.collected_trace( - tenant_id=tenant_id, app_id=link.app_id, workflow_run_id=link.workflow_run_id - ) - if self._collector is not None - else None - ) - if collected is None: - trace_id = root_trace_id(link.external_trace_id, root.root_id) - external_trace_id = link.external_trace_id - outer_destination = self.destination(tenant_id, link.app_id) - elif collected.trace_id is None: - return - else: - # An Agent trace keeps a caller's trace ID only as an attribute. - trace_id, external_trace_id = collected.trace_id, None - outer_destination = ( - self._decoded(tenant_id, link.app_id, collected.destination) - if collected.destination is not None - else None - ) - yield OuterTrace( - trace_id=trace_id, - external_trace_id=external_trace_id, - root=SpanRef(span_id=span_id_for(root.root_id), start_time=root.started_at), - same_destination=outer_destination is not None and outer_destination.is_same(destination), - collected=collected is not None, - ) - link = link.parent - - def _decoded(self, tenant_id: str, app_id: str, stored: StoredTraceDestination) -> Destination | None: + if stored is None: + return None try: provider = self._providers.resolve(stored.provider_id) config = self._codec.decode(tenant_id, provider, stored.config) diff --git a/api/services/tracing/records.py b/api/services/tracing/records.py index 448ca0e882e1d6..0e1a911a916304 100644 --- a/api/services/tracing/records.py +++ b/api/services/tracing/records.py @@ -11,7 +11,6 @@ from enum import StrEnum from typing import Any -from dify_trace_sdk.events import TraceParent from pydantic import BaseModel, ConfigDict, StrictStr, ValidationError from graphon.entities import WorkflowNodeExecution @@ -115,8 +114,6 @@ class WorkflowTraceInfo(BaseTraceInfo): file_list: list[str] invoked_by: str | None query: str - parent: TraceParent | None - """The run that started this nested workflow run; ``None`` outside nested workflows.""" class MessageTraceInfo(BaseTraceInfo): diff --git a/api/services/tracing/recovery.py b/api/services/tracing/recovery.py index a3fa548b0f7cd6..18e69ff0c16234 100644 --- a/api/services/tracing/recovery.py +++ b/api/services/tracing/recovery.py @@ -39,7 +39,8 @@ def _event(run: TraceRunAgg, payload: TracePayload, user_id: str | None) -> Trac correlation=TraceCorrelation( external_trace_id=None, session_id=run.session_id, - parent=None, + parent_workflow_run_id=None, + parent_node_execution_id=None, ), payload=payload, ) diff --git a/api/services/tracing/scopes.py b/api/services/tracing/scopes.py index 1652eafa32955b..86f3a2edbe9f1f 100644 --- a/api/services/tracing/scopes.py +++ b/api/services/tracing/scopes.py @@ -16,7 +16,9 @@ logger = logging.getLogger(__name__) -UNCORRELATED = TraceCorrelation(external_trace_id=None, session_id=None, parent=None) +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.""" diff --git a/api/services/tracing/service.py b/api/services/tracing/service.py index 25a2d267262cf4..39c1a8fcfd827a 100644 --- a/api/services/tracing/service.py +++ b/api/services/tracing/service.py @@ -40,7 +40,7 @@ workflow_has_agent_v2, ) from services.tracing.collection_ports import ScheduleTraceExport, TraceRunQuery, TraceRunRepository, TraceSourceQuery -from services.tracing.delivery import CollectedTrace, StoredTraceDestination, TraceCollector, TraceRecordSource +from services.tracing.delivery import TraceCollector, TraceRecordSource from services.tracing.errors import ( ProviderConfigUnreadableError, TraceCollectionConflictError, @@ -337,7 +337,8 @@ def capture_agent_app_terminal( correlation=TraceCorrelation( external_trace_id=None, session_id=run.session_id if run else source.conversation_id, - parent=None, + parent_workflow_run_id=None, + parent_node_execution_id=None, ), payload=MessageCompleted(message_id=message_id, conversation_id=source.conversation_id), ) @@ -431,20 +432,24 @@ def _managed_owner( return False @override - def collected_trace(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> CollectedTrace | None: - run = self._run_query.workflow_run_trace(tenant_id=tenant_id, workflow_run_id=workflow_run_id) - if run is not None and run.app_id != app_id: - # A collected run of another app is not the run that started the nested run. - return CollectedTrace(trace_id=None, destination=None) - if run is not None: - destination = ( - StoredTraceDestination(provider_id=run.provider, config=run.provider_config) if run.provider else None - ) - return CollectedTrace(trace_id=run.trace_id, destination=destination) - app = self._sources.app(tenant_id=tenant_id, app_id=app_id) - if app is None or not self._managed_owner(app, None, workflow_run_id): - return None - return CollectedTrace(trace_id=None, destination=None) + 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 info + parent = self._run_query.workflow_run_trace(tenant_id=tenant_id, workflow_run_id=parent_run_id) + if parent is None: + 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 + if parent_node_id: + 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 + return info.model_copy(update={"metadata": metadata}) @override def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: 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 83eb9a52c21fff..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 @@ -66,7 +66,9 @@ def _trace_scope() -> TraceScope: tenant_id="tenant", app_id="app", user_id="user", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) 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 e56c452c67aefc..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 @@ -100,7 +100,9 @@ def _trace_scope() -> TraceScope: tenant_id="tenant", app_id="app", user_id="user", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) 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 6a2cd410032cf8..9f2f0bafc60747 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 @@ -93,7 +93,9 @@ def _trace_scope() -> TraceScope: tenant_id="tenant", app_id="app", user_id="user-id", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) @@ -1664,7 +1666,9 @@ def raise_conversation_not_exists(**_kwargs): tenant_id="tenant-id", app_id="app-id", user_id="user-id", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + 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]] = [] 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 bf7ba67f3f5c57..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 @@ -95,7 +95,8 @@ def _trace_scope(recorder: _TraceRecorder) -> TraceScope: correlation=TraceCorrelation( external_trace_id=None, session_id="trace-session-1", - parent=None, + parent_workflow_run_id=None, + parent_node_execution_id=None, ), ) @@ -436,7 +437,9 @@ def _timeline_trace_scope(recorder: _TimelineTraceRecorder) -> TraceScope: tenant_id="tenant", app_id="app", user_id="session", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) 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 29658525afe6e5..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 @@ -85,7 +85,9 @@ def _trace_scope() -> TraceScope: tenant_id="tenant", app_id="app1", user_id="user", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) @@ -611,7 +613,9 @@ def _make_entity(query: str) -> AgentAppGenerateEntity: tenant_id="tenant-1", app_id="app-1", user_id="user-1", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + 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", 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 bbee5286394885..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 @@ -60,7 +60,9 @@ def _trace_scope() -> TraceScope: tenant_id="tenant", app_id="app1", user_id="user", - correlation=TraceCorrelation(external_trace_id=None, session_id="session-1", parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id="session-1", parent_workflow_run_id=None, parent_node_execution_id=None + ), ) 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 633efafaaae3b7..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 @@ -40,7 +40,9 @@ def _trace_scope() -> TraceScope: tenant_id="tenant-1", app_id="app-1", user_id="user-1", - correlation=TraceCorrelation(external_trace_id=None, session_id="session-1", parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id="session-1", parent_workflow_run_id=None, parent_node_execution_id=None + ), ) 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 b09292b25b09f1..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 @@ -36,7 +36,9 @@ def _trace_scope() -> TraceScope: tenant_id="tenant", app_id="app1", user_id="user", - correlation=TraceCorrelation(external_trace_id=None, session_id="session-1", parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id="session-1", parent_workflow_run_id=None, parent_node_execution_id=None + ), ) 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 0a4021ee5b6156..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 @@ -889,7 +889,9 @@ def check(self: object, **kwargs: object) -> tuple[bool, dict[str, object], str] tenant_id="tenant", app_id="app", user_id="user", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + 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", diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py index 60a9a4f94755ae..4ffc0bde565467 100644 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py @@ -48,7 +48,6 @@ ) from core.credit_usage import CreditUsageAppType from core.db.session_factory import session_factory -from core.helper.trace_id_helper import ParentTraceContext from core.repositories.human_input_repository import HumanInputFormRepositoryImpl, HumanInputFormSubmissionRepository from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository @@ -345,19 +344,10 @@ def fake_from_graph_init_context(**kwargs): invoke_from=InvokeFrom.DEBUGGER, root_node_id="root", trace_session_id="session-1", - external_trace_id="caller-trace", - parent_trace_context=ParentTraceContext( - parent_workflow_run_id="outer-run", parent_node_execution_id="outer-node-execution" - ), call_depth=3, ) - dify_context = captured["run_context"][DIFY_RUN_CONTEXT_KEY] - assert dify_context.trace_session_id == "session-1" - assert dify_context.external_trace_id == "caller-trace" - assert dify_context.parent_trace_context == ParentTraceContext( - parent_workflow_run_id="outer-run", parent_node_execution_id="outer-node-execution" - ) + assert captured["run_context"][DIFY_RUN_CONTEXT_KEY].trace_session_id == "session-1" assert captured["call_depth"] == 3 assert captured["runtime_state"] is runtime_state @@ -1190,8 +1180,6 @@ def test_run_uses_single_node_execution_branch( user_id="user", app_type=CreditUsageAppType.WORKFLOW, trace_session_id="session-1", - external_trace_id=None, - parent_trace_context=None, ) init_graph.assert_not_called() 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 2bea85a00aad14..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 @@ -31,7 +31,6 @@ _WorkflowGenerateEntityWrapper, ) from core.app.workflow.file_runtime import DifyWorkflowFileRuntime -from core.helper.trace_id_helper import ParentTraceContext from core.repositories import SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from core.tools.tool_file_manager import ToolFileManager @@ -84,7 +83,9 @@ def _trace_scope(recorder: _TraceRecorder | _PersistedRunRecorder, *, app: App, 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=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) @@ -398,9 +399,10 @@ def fake_generate(**kwargs): assert isinstance(application_generate_entity, WorkflowAppGenerateEntity) extras = application_generate_entity.extras assert extras["external_trace_id"] == "trace-1" - assert extras["parent_trace_context"] == ParentTraceContext( - parent_workflow_run_id="outer-workflow-run-1", parent_node_execution_id="outer-node-execution-1" - ) + assert extras["parent_trace_context"].model_dump() == { + "parent_workflow_run_id": "outer-workflow-run-1", + "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 == [ @@ -716,7 +718,9 @@ def _app_generator_extra_trace_scope() -> TraceScope: tenant_id=TENANT_ID, app_id=APP_ID, user_id="session-id", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) 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 39af65e7bd9cba..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 @@ -40,7 +40,9 @@ def _trace_scope() -> TraceScope: tenant_id="tenant-id", app_id="app-id", user_id="user-id", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) 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 94fc4183a358d1..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 @@ -455,7 +455,9 @@ def test_message_end_emits_message_trace_after_commit( 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=None), + 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] = [] @@ -560,7 +562,9 @@ def _trace_scope(recorder: _TimelineTraceRecorder) -> TraceScope: tenant_id="tenant", app_id="app", user_id="user", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) 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 a0893ac1c56d68..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 @@ -215,7 +215,8 @@ def test_generate_conversation_name_success( correlation=TraceCorrelation( external_trace_id="caller-trace", session_id=None, - parent=None, + parent_workflow_run_id=None, + parent_node_execution_id=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 cf7f51c0b7b2c5..ad51d7bbc1489d 100644 --- a/api/tests/unit_tests/core/moderation/test_input_moderation.py +++ b/api/tests/unit_tests/core/moderation/test_input_moderation.py @@ -26,7 +26,9 @@ def _trace_scope(recorder: _Recorder) -> TraceScope: 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=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) 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 d7655a6e1c9bd7..2dfb989a6367a9 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 @@ -98,7 +98,9 @@ def _trace_scope(recorder: _Recorder) -> TraceScope: tenant_id="tenant", app_id="app", user_id="user", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) diff --git a/api/tests/unit_tests/core/telemetry/test___init__.py b/api/tests/unit_tests/core/telemetry/test___init__.py index af6799d0a5093f..2692242d087127 100644 --- a/api/tests/unit_tests/core/telemetry/test___init__.py +++ b/api/tests/unit_tests/core/telemetry/test___init__.py @@ -33,7 +33,9 @@ def create_scope( tenant_id=tenant_id, app_id=app_id, user_id=user_id, - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + 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: diff --git a/api/tests/unit_tests/core/telemetry/test_gateway.py b/api/tests/unit_tests/core/telemetry/test_gateway.py index 4b39504ff353f9..6df19f851f74a3 100644 --- a/api/tests/unit_tests/core/telemetry/test_gateway.py +++ b/api/tests/unit_tests/core/telemetry/test_gateway.py @@ -120,7 +120,9 @@ def create_scope( tenant_id=tenant_id, app_id=app_id, user_id=user_id, - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + 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: 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 3ead7cdad779c3..54abd16923187e 100644 --- a/api/tests/unit_tests/core/tools/test_tool_engine.py +++ b/api/tests/unit_tests/core/tools/test_tool_engine.py @@ -57,7 +57,9 @@ def _trace_scope(recorder: _Recorder) -> TraceScope: tenant_id="tenant-1", app_id="app-1", user_id="u1", - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) diff --git a/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py b/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py index 109639d4776e81..d066c8fe771fb5 100644 --- a/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py +++ b/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py @@ -13,7 +13,6 @@ from sqlalchemy.orm import Session, sessionmaker from core.app.entities.app_invoke_entities import InvokeFrom -from core.helper.trace_id_helper import ParentTraceContext from core.tools.__base.tool_runtime import ToolRuntime from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ( @@ -340,15 +339,8 @@ def test_workflow_tool_passes_parent_trace_context_from_runtime( """Ensure nested workflow runtime metadata is forwarded as parent trace context.""" tool = _build_tool() tool.set_parent_trace_context( - ParentTraceContext( - parent_workflow_run_id="outer-workflow-run-1", - parent_node_execution_id="outer-node-execution-1", - parent_trace_context=ParentTraceContext( - parent_workflow_run_id="outermost-workflow-run", - parent_node_execution_id="outermost-node-execution", - parent_external_trace_id="caller-trace", - ), - ) + parent_workflow_run_id="outer-workflow-run-1", + parent_node_execution_id="outer-node-execution-1", ) monkeypatch.setattr(tool, "_get_app", lambda *args, **kwargs: None) @@ -366,15 +358,6 @@ def test_workflow_tool_passes_parent_trace_context_from_runtime( assert call_kwargs["args"]["parent_trace_context"].model_dump() == { "parent_workflow_run_id": "outer-workflow-run-1", "parent_node_execution_id": "outer-node-execution-1", - "parent_app_id": None, - "parent_external_trace_id": None, - "parent_trace_context": { - "parent_workflow_run_id": "outermost-workflow-run", - "parent_node_execution_id": "outermost-node-execution", - "parent_app_id": None, - "parent_external_trace_id": "caller-trace", - "parent_trace_context": None, - }, } @@ -431,9 +414,8 @@ def test_workflow_tool_keeps_user_inputs_named_like_trace_runtime_keys( ), ] tool.set_parent_trace_context( - ParentTraceContext( - parent_workflow_run_id="outer-workflow-run-1", parent_node_execution_id="outer-node-execution-1" - ) + parent_workflow_run_id="outer-workflow-run-1", + parent_node_execution_id="outer-node-execution-1", ) monkeypatch.setattr(tool, "_get_app", lambda *args, **kwargs: None) @@ -459,9 +441,10 @@ def test_workflow_tool_keeps_user_inputs_named_like_trace_runtime_keys( call_kwargs = generate_calls[-1] assert call_kwargs["args"]["inputs"]["outer_workflow_run_id"] == "user-workflow-input" assert call_kwargs["args"]["inputs"]["outer_node_execution_id"] == "user-node-input" - assert call_kwargs["args"]["parent_trace_context"] == ParentTraceContext( - parent_workflow_run_id="outer-workflow-run-1", parent_node_execution_id="outer-node-execution-1" - ) + assert call_kwargs["args"]["parent_trace_context"].model_dump() == { + "parent_workflow_run_id": "outer-workflow-run-1", + "parent_node_execution_id": "outer-node-execution-1", + } def test_workflow_tool_can_clear_parent_trace_context( @@ -471,9 +454,8 @@ def test_workflow_tool_can_clear_parent_trace_context( """Ensure reused WorkflowTool instances do not keep stale parent trace context.""" tool = _build_tool() tool.set_parent_trace_context( - ParentTraceContext( - parent_workflow_run_id="outer-workflow-run-1", parent_node_execution_id="outer-node-execution-1" - ) + parent_workflow_run_id="outer-workflow-run-1", + parent_node_execution_id="outer-node-execution-1", ) tool.clear_parent_trace_context() 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 34d3075ccde8b4..e541b765b8df20 100644 --- a/api/tests/unit_tests/core/workflow/test_node_runtime.py +++ b/api/tests/unit_tests/core/workflow/test_node_runtime.py @@ -21,7 +21,6 @@ from core.app.llm.model_access import DifyCredentialsProvider, DifyModelFactory from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler from core.helper.ssrf_proxy import graphon_ssrf_proxy -from core.helper.trace_id_helper import ParentTraceContext from core.llm_generator.output_parser.errors import OutputParserError from core.llm_generator.output_parser.structured_output import _parse_structured_output from core.model_manager import QuotaManagedModelInstance @@ -1143,7 +1142,7 @@ def test_get_runtime_converts_graph_provider_type_for_tool_manager(runtime: Dify assert workflow_tool.provider_type == CoreToolProviderType.BUILT_IN @staticmethod - def test_workflow_tools_receive_the_chained_parent_trace_context( + def test_get_runtime_stores_parent_trace_context_for_workflow_tools( runtime: DifyToolNodeRuntime, ) -> None: variable_pool: VariablePool = build_test_variable_pool( @@ -1154,11 +1153,6 @@ def test_workflow_tools_receive_the_chained_parent_trace_context( ) workflow_runtime = MagicMock() workflow_runtime.runtime.runtime_parameters = {} - runtime._run_context.parent_trace_context = ParentTraceContext( - parent_workflow_run_id="outer-run-id", - parent_node_execution_id="outer-node-execution-id", - parent_external_trace_id="caller-trace", - ) node_data = ToolNodeData.model_validate( { "type": "tool", @@ -1181,37 +1175,14 @@ def test_workflow_tools_receive_the_chained_parent_trace_context( node_execution_id="node-execution-id", ) - chained = ParentTraceContext( - parent_workflow_run_id="workflow-run-id", - parent_node_execution_id="node-execution-id", - parent_app_id="app-id", - parent_trace_context=ParentTraceContext( - parent_workflow_run_id="outer-run-id", - parent_node_execution_id="outer-node-execution-id", - parent_external_trace_id="caller-trace", - ), - ) - assert tool_runtime.raw.parent_trace_context == chained + assert tool_runtime.raw.parent_trace_context.model_dump() == { + "parent_workflow_run_id": "workflow-run-id", + "parent_node_execution_id": "node-execution-id", + } assert workflow_runtime.runtime.runtime_parameters == {} - with ( - patch.object(ToolEngine, "generic_invoke", return_value=iter([])), - patch.object( - ToolFileMessageTransformer, - "transform_tool_invoke_messages", - side_effect=lambda *, messages, **_: messages, - ), - ): - list( - runtime.invoke( - tool_runtime=tool_runtime, tool_parameters={}, workflow_call_depth=0, provider_name="provider" - ) - ) - - workflow_runtime.set_parent_trace_context.assert_called_once_with(chained) - @staticmethod - def test_get_runtime_passes_trace_session_and_caller_trace_to_workflow_tools( + def test_get_runtime_stores_trace_session_id_for_workflow_tools( runtime: DifyToolNodeRuntime, ) -> None: variable_pool: VariablePool = build_test_variable_pool( @@ -1223,7 +1194,6 @@ def test_get_runtime_passes_trace_session_and_caller_trace_to_workflow_tools( workflow_runtime = MagicMock() workflow_runtime.runtime.runtime_parameters = {} runtime._run_context.trace_session_id = "session-1" - runtime._run_context.external_trace_id = "caller-trace" node_data = ToolNodeData.model_validate( { "type": "tool", @@ -1247,12 +1217,6 @@ def test_get_runtime_passes_trace_session_and_caller_trace_to_workflow_tools( ) assert tool_runtime.raw.trace_session_id == "session-1" - assert tool_runtime.raw.parent_trace_context == ParentTraceContext( - parent_workflow_run_id="workflow-run-id", - parent_node_execution_id="node-execution-id", - parent_app_id="app-id", - parent_external_trace_id="caller-trace", - ) assert workflow_runtime.runtime.runtime_parameters == {} @staticmethod @@ -2210,11 +2174,10 @@ def test_dify_tool_node_runtime_injects_outer_workflow_run_id_for_workflow_tools ) assert handle.raw.tool is runtime_tool - assert handle.raw.parent_trace_context == ParentTraceContext( - parent_workflow_run_id="outer-workflow-run-id", - parent_node_execution_id="node-execution-id", - parent_app_id="app-id", - ) + assert handle.raw.parent_trace_context.model_dump() == { + "parent_workflow_run_id": "outer-workflow-run-id", + "parent_node_execution_id": "node-execution-id", + } assert runtime_tool.runtime.runtime_parameters == {} get_runtime.assert_called_once() diff --git a/api/tests/unit_tests/domain/tracing/test_nested_workflow_trace.py b/api/tests/unit_tests/domain/tracing/test_nested_workflow_trace.py deleted file mode 100644 index 4dd2e804e639bb..00000000000000 --- a/api/tests/unit_tests/domain/tracing/test_nested_workflow_trace.py +++ /dev/null @@ -1,108 +0,0 @@ -from collections.abc import Iterator -from datetime import UTC, datetime, timedelta - -from dify_trace_sdk.ids import span_id_for, trace_id_for -from dify_trace_sdk.model import Span, SpanKind, SpanRef, SpanStatus, Trace - -from domain.tracing.nested_workflow_trace import OuterTrace, TracePlacement, nested_trace_placement - -STARTED_AT = datetime(2026, 10, 1, 12, 0, 0, tzinfo=UTC) -TOOL_NODE = SpanRef(span_id=span_id_for("outer-tool-node"), start_time=STARTED_AT + timedelta(seconds=2)) - - -def outer(run: str, *, same_destination: bool = True, collected: bool = False) -> OuterTrace: - return OuterTrace( - trace_id=trace_id_for(run), - external_trace_id=f"{run}-caller", - root=SpanRef(span_id=span_id_for(run), start_time=STARTED_AT), - same_destination=same_destination, - collected=collected, - ) - - -def joined(run: str) -> TracePlacement: - return TracePlacement( - trace_id=trace_id_for(run), - external_trace_id=f"{run}-caller", - parent=TOOL_NODE, - trace_root=SpanRef(span_id=span_id_for(run), start_time=STARTED_AT), - ) - - -def span(span_id: str, parent_span_id: str | None) -> Span: - return Span( - span_id=span_id, - parent_span_id=parent_span_id, - name="workflow", - kind=SpanKind.CHAIN, - start_time=STARTED_AT + timedelta(seconds=3), - end_time=STARTED_AT + timedelta(seconds=4), - status=SpanStatus.OK, - ) - - -class TestNestedTracePlacement: - @staticmethod - def test_nested_run_hangs_from_the_tool_node_in_the_trace_of_its_outer_run() -> None: - assert nested_trace_placement([outer("parent")], tool_node=TOOL_NODE) == joined("parent") - - @staticmethod - def test_outer_runs_sharing_the_destination_join_the_outermost_trace() -> None: - placement = nested_trace_placement([outer("parent"), outer("grandparent")], tool_node=TOOL_NODE) - - assert placement == joined("grandparent") - - @staticmethod - def test_an_outer_run_exporting_elsewhere_starts_the_trace_below_it() -> None: - assert nested_trace_placement( - [outer("parent"), outer("grandparent", same_destination=False)], tool_node=TOOL_NODE - ) == joined("parent") - assert ( - nested_trace_placement([outer("parent", same_destination=False), outer("grandparent")], tool_node=TOOL_NODE) - is None - ) - - @staticmethod - def test_a_collected_outer_run_is_the_root_of_its_trace() -> None: - placement = nested_trace_placement([outer("parent", collected=True), outer("grandparent")], tool_node=TOOL_NODE) - - assert placement == joined("parent") - - @staticmethod - def test_a_nested_run_without_a_known_outer_trace_starts_its_own() -> None: - assert nested_trace_placement([], tool_node=TOOL_NODE) is None - - @staticmethod - def test_outer_runs_are_read_only_as_far_as_the_decision_needs() -> None: - def outer_traces() -> Iterator[OuterTrace]: - yield outer("parent", same_destination=False) - raise AssertionError("the run above an outer run exporting elsewhere was read") - - assert nested_trace_placement(outer_traces(), tool_node=TOOL_NODE) is None - - -class TestTracePlacement: - @staticmethod - def test_placed_trace_hangs_its_root_from_the_tool_node_of_the_outer_trace() -> None: - root_span_id, node_span_id = span_id_for("nested-run"), span_id_for("nested-node") - trace = Trace( - trace_id=trace_id_for("nested-run"), - external_trace_id=None, - session_id="session", - user_id="account", - parent=None, - trace_root=None, - spans=(span(root_span_id, None), span(node_span_id, root_span_id)), - ) - - placed = joined("parent").place(trace) - - assert placed.trace_id == trace_id_for("parent") - assert placed.external_trace_id == "parent-caller" - assert (placed.session_id, placed.user_id) == ("session", "account") - assert placed.parent == TOOL_NODE - assert placed.trace_root == SpanRef(span_id=span_id_for("parent"), start_time=STARTED_AT) - assert [(item.span_id, item.parent_span_id) for item in placed.spans] == [ - (root_span_id, TOOL_NODE.span_id), - (node_span_id, root_span_id), - ] diff --git a/api/tests/unit_tests/domain/tracing/test_normalization.py b/api/tests/unit_tests/domain/tracing/test_normalization.py index b2477593821ea3..fa74baab3ad8a4 100644 --- a/api/tests/unit_tests/domain/tracing/test_normalization.py +++ b/api/tests/unit_tests/domain/tracing/test_normalization.py @@ -179,7 +179,6 @@ def trace_of(*spans: Span) -> Trace: session_id="session", user_id="user", parent=None, - trace_root=None, spans=spans, ) @@ -273,7 +272,6 @@ def test_managed_external_parent_is_preserved_as_explicit_link_when_content_off( session_id="session", user_id=None, parent=None, - trace_root=None, spans=(root,), ) remapped = remap_outer_trace(outer, trace_id="1" * 32, capture_content=False, synthetic_span_ids=frozenset()) 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 index 6775db0123b354..8fe1b1d90b659c 100644 --- a/api/tests/unit_tests/domain/tracing/test_trace_run_agg.py +++ b/api/tests/unit_tests/domain/tracing/test_trace_run_agg.py @@ -95,7 +95,6 @@ def business_terminal(run: TraceRunAgg) -> Trace: session_id="conversation", user_id="user", parent=None, - trace_root=None, spans=( Span( span_id=span_id_for(run.source_id), 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 05dccbf8854690..90502db2cca719 100644 --- a/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py +++ b/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py @@ -334,7 +334,9 @@ def create_scope( tenant_id=tenant_id, app_id=app_id, user_id=user_id, - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + 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: 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 b730c487ff876b..5aa4c413cefc6e 100644 --- a/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py +++ b/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py @@ -98,7 +98,6 @@ def make_workflow_info(**overrides) -> WorkflowTraceInfo: "completion_tokens": 40, "file_list": [], "query": "hello", - "parent": None, "start_time": _T0, "end_time": _T1, "metadata": { diff --git a/api/tests/unit_tests/extensions/test_tracing.py b/api/tests/unit_tests/extensions/test_tracing.py index 446ac81843a062..64ee89e7f73309 100644 --- a/api/tests/unit_tests/extensions/test_tracing.py +++ b/api/tests/unit_tests/extensions/test_tracing.py @@ -2,7 +2,7 @@ from dify_agent.layers.execution_context import ApplicationTracingContext from dify_trace_sdk.emitter import TraceScope -from dify_trace_sdk.events import TraceCorrelation, TraceEvent, TraceParent +from dify_trace_sdk.events import TraceCorrelation, TraceEvent from flask import Flask from core.helper.trace_id_helper import ParentTraceContext @@ -113,17 +113,7 @@ def test_correlation_reads_the_caller_extras() -> None: { "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", - "parent_app_id": "outer-app", - "parent_trace_context": { - "parent_workflow_run_id": "outermost-run", - "parent_node_execution_id": "outermost-node", - "parent_app_id": "outermost-app", - "parent_external_trace_id": "outermost-caller-trace", - }, - }, + "parent_trace_context": {"parent_workflow_run_id": "outer-run", "parent_node_execution_id": "outer-node"}, } ) typed = correlation_from_extras( @@ -133,25 +123,12 @@ def test_correlation_reads_the_caller_extras() -> None: assert nested == TraceCorrelation( external_trace_id="caller-trace", session_id="session-1", - parent=TraceParent( - workflow_run_id="outer-run", - app_id="outer-app", - node_execution_id="outer-node", - external_trace_id=None, - parent=TraceParent( - workflow_run_id="outermost-run", - app_id="outermost-app", - node_execution_id="outermost-node", - external_trace_id="outermost-caller-trace", - parent=None, - ), - ), - ) - assert typed.parent == TraceParent( - workflow_run_id="outer-run", app_id=None, node_execution_id=None, external_trace_id=None, parent=None + 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=None + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=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 a8345389423d20..2cc696c336ba2a 100644 --- a/api/tests/unit_tests/repositories/test_application_trace_repository.py +++ b/api/tests/unit_tests/repositories/test_application_trace_repository.py @@ -36,7 +36,7 @@ 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 CollectedTrace, StoredTraceDestination, TraceRecordSource +from services.tracing.delivery import TraceRecordSource from services.tracing.errors import ( TraceCollectionConflictError, TraceCollectionError, @@ -224,7 +224,6 @@ def agent_turn(run: TraceRun, *, root_span_id: str) -> Trace: session_id=run.session_id, user_id=None, parent=None, - trace_root=None, spans=(root,), ) @@ -333,7 +332,6 @@ def workflow_info(row: WorkflowRun, parent: WorkflowRun | None) -> WorkflowTrace file_list=[], invoked_by=None, query="", - parent=None, metadata=metadata, start_time=row.created_at, end_time=row.finished_at, @@ -723,22 +721,21 @@ def test_auxiliary_content_is_suppressed_for_logstore_chatflow(stored: StoredSou @staticmethod @pytest.mark.parametrize("stored", ["workflow"], indirect=True) - def test_collected_trace_is_the_frozen_agent_trace_of_the_runs_own_app(stored: StoredSource) -> None: + @pytest.mark.parametrize("mismatch", [None, "tenant_id", "app_id"]) + def test_linked_parent_requires_logstore_run_and_durable_owner_to_agree( + stored: StoredSource, mismatch: str | None + ) -> None: parent = stored.start() - service = stored.service() - - collected = service.collected_trace( - tenant_id=parent.tenant_id, app_id=parent.app_id, workflow_run_id=parent.source_id - ) - foreign = service.collected_trace( - tenant_id=parent.tenant_id, app_id=identifier("foreign"), workflow_run_id=parent.source_id - ) - - assert collected == CollectedTrace( - trace_id=parent.trace_id, - destination=StoredTraceDestination(provider_id=parent.provider, config=parent.provider_config), - ) - assert foreign == CollectedTrace(trace_id=None, destination=None) + info = workflow_info(_dict_to_workflow_run(stored.row), _dict_to_workflow_run(stored.row)) + if mismatch: + stored.row[mismatch] = identifier("foreign") + linked = stored.service().link_managed_parent(tenant_id=stored.case.run.tenant_id, info=info) + if mismatch: + assert linked is info + else: + 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( diff --git a/api/tests/unit_tests/repositories/test_trace_record_repository.py b/api/tests/unit_tests/repositories/test_trace_record_repository.py index 8b2957575ee27b..51f882d3fdaa62 100644 --- a/api/tests/unit_tests/repositories/test_trace_record_repository.py +++ b/api/tests/unit_tests/repositories/test_trace_record_repository.py @@ -13,24 +13,9 @@ from sqlalchemy import create_engine from sqlalchemy.orm import Session, sessionmaker -from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus -from models.enums import ( - ConversationFromSource, - ConversationStatus, - CreatorUserRole, - EndUserType, - InvokeFrom, - WorkflowRunTriggeredFrom, -) +from models.enums import ConversationFromSource, ConversationStatus, EndUserType, InvokeFrom from models.model import App, AppMode, Conversation, EndUser, Message -from models.workflow import ( - WorkflowNodeExecutionModel, - WorkflowNodeExecutionTriggeredFrom, - WorkflowRun, - WorkflowType, -) from repositories.trace_record_repository import TraceRecordRepository -from services.tracing.delivery import OuterRunRoot from services.tracing.errors import SourceNotFoundError, SourceNotReadyError, TracingStoreUnavailableError from services.tracing.records import MessageTraceInfo, ModerationTraceInfo, ToolTraceInfo, TraceMessage @@ -38,10 +23,6 @@ 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" -CHATFLOW_RUN_ID = "outer-chatflow-run" -TOOL_NODE_ID = "outer-tool-node" @pytest.fixture @@ -110,77 +91,6 @@ def records(sqlite_session: Session) -> None: 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 running workflow run whose tool node starts nested runs, and the chatflow run the message answers.""" - 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) - chat_app = sqlite_session.get(App, APP_ID) - message = sqlite_session.get(Message, MESSAGE_ID) - assert chat_app is not None - assert message is not None - chat_app.mode = AppMode.ADVANCED_CHAT - message.workflow_run_id = CHATFLOW_RUN_ID - sqlite_session.add_all( - [workflow_app, _workflow_run(OUTER_RUN_ID, WORKFLOW_APP_ID), tool_node, _workflow_run(CHATFLOW_RUN_ID, APP_ID)] - ) - sqlite_session.commit() - - def _assembler(sessions: sessionmaker[Session]) -> TraceRecordRepository: return TraceRecordRepository(sessions, files_url="https://files.example.com", enterprise_details=False) @@ -195,7 +105,8 @@ def _event(payload: TracePayload, *, tenant_id: str = TENANT_ID) -> TraceEvent: correlation=TraceCorrelation( external_trace_id="caller-trace", session_id="session-1", - parent=None, + parent_workflow_run_id=None, + parent_node_execution_id=None, ), payload=payload, ) @@ -324,51 +235,3 @@ def test_an_unreachable_database_is_reported_as_unavailable() -> None: with pytest.raises(TracingStoreUnavailableError): _assembler(sessions).assemble(_event(MessageCompleted(message_id=MESSAGE_ID, conversation_id="conversation-1"))) - - -@pytest.mark.usefixtures("records", "outer_runs") -def test_a_workflow_trace_starts_with_its_run_and_a_chatflow_trace_with_its_message( - sqlite_session_factory: sessionmaker[Session], -) -> None: - repository = _assembler(sqlite_session_factory) - - assert repository.root(tenant_id=TENANT_ID, app_id=WORKFLOW_APP_ID, workflow_run_id=OUTER_RUN_ID) == OuterRunRoot( - root_id=OUTER_RUN_ID, started_at=CREATED_AT + timedelta(seconds=1) - ) - assert repository.root(tenant_id=TENANT_ID, app_id=APP_ID, workflow_run_id=CHATFLOW_RUN_ID) == OuterRunRoot( - root_id=MESSAGE_ID, started_at=CREATED_AT - ) - - -@pytest.mark.usefixtures("records", "outer_runs") -def test_outer_runs_and_their_tool_nodes_are_read_through_their_owners( - sqlite_session_factory: sessionmaker[Session], -) -> None: - repository = _assembler(sqlite_session_factory) - - def started_at(tenant_id: str, app_id: str, workflow_run_id: str) -> datetime | None: - return repository.node_started_at( - tenant_id=tenant_id, app_id=app_id, workflow_run_id=workflow_run_id, node_execution_id=TOOL_NODE_ID - ) - - assert started_at(TENANT_ID, WORKFLOW_APP_ID, OUTER_RUN_ID) == CREATED_AT + timedelta(seconds=2) - assert started_at("other-tenant", WORKFLOW_APP_ID, OUTER_RUN_ID) is None - assert started_at(TENANT_ID, APP_ID, OUTER_RUN_ID) is None - assert started_at(TENANT_ID, WORKFLOW_APP_ID, CHATFLOW_RUN_ID) is None - assert repository.root(tenant_id="other-tenant", app_id=WORKFLOW_APP_ID, workflow_run_id=OUTER_RUN_ID) is None - assert repository.root(tenant_id=TENANT_ID, app_id=APP_ID, workflow_run_id=OUTER_RUN_ID) is None - - -@pytest.mark.usefixtures("records", "outer_runs") -def test_a_chatflow_run_has_no_root_until_its_message_names_it( - sqlite_session_factory: sessionmaker[Session], sqlite_session: Session -) -> None: - message = sqlite_session.get(Message, MESSAGE_ID) - assert message is not None - message.workflow_run_id = None - sqlite_session.commit() - - assert ( - _assembler(sqlite_session_factory).root(tenant_id=TENANT_ID, app_id=APP_ID, workflow_run_id=CHATFLOW_RUN_ID) - is None - ) diff --git a/api/tests/unit_tests/services/test_workflow_persistence.py b/api/tests/unit_tests/services/test_workflow_persistence.py index ad3f839319fc10..ea38df25573781 100644 --- a/api/tests/unit_tests/services/test_workflow_persistence.py +++ b/api/tests/unit_tests/services/test_workflow_persistence.py @@ -8,7 +8,7 @@ import pytest from dify_trace_sdk.emitter import TraceScope -from dify_trace_sdk.events import TraceCorrelation, TraceEvent, TraceParent, WorkflowRunCompleted +from dify_trace_sdk.events import TraceCorrelation, TraceEvent, WorkflowRunCompleted from flask import Flask from sqlalchemy import event as sqlalchemy_event from sqlalchemy import select @@ -293,7 +293,9 @@ def _trace_scope(recorder: _TraceRecorder, correlation: TraceCorrelation) -> Tra def _empty_correlation() -> TraceCorrelation: - return TraceCorrelation(external_trace_id=None, session_id=None, parent=None) + return TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ) def _naive_utc_now() -> datetime: @@ -1102,13 +1104,8 @@ def test_handle_graph_run_succeeded_emits_run_tokens_with_scope_correlation( correlation = TraceCorrelation( external_trace_id="trace", session_id="session-1", - parent=TraceParent( - workflow_run_id="outer-workflow-run-1", - app_id="outer-app-1", - node_execution_id="outer-node-execution-1", - external_trace_id=None, - parent=None, - ), + parent_workflow_run_id="outer-workflow-run-1", + parent_node_execution_id="outer-node-execution-1", ) layer, _, _, runtime_state = _make_layer( system_variables, diff --git a/api/tests/unit_tests/services/tracing/test_adapters.py b/api/tests/unit_tests/services/tracing/test_adapters.py index ff46c3a61805cb..5afba1f01b2a70 100644 --- a/api/tests/unit_tests/services/tracing/test_adapters.py +++ b/api/tests/unit_tests/services/tracing/test_adapters.py @@ -24,7 +24,9 @@ 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=None) +CORRELATION = TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None +) class _Broker: @@ -150,7 +152,6 @@ def test_export_retries_persist_the_trace_and_run_after_their_delay(broker: _Bro session_id=None, user_id=None, parent=None, - trace_root=None, spans=( Span( span_id="b" * 16, diff --git a/api/tests/unit_tests/services/tracing/test_canonical.py b/api/tests/unit_tests/services/tracing/test_canonical.py index de91f7b270d76a..1c22b6f305b150 100644 --- a/api/tests/unit_tests/services/tracing/test_canonical.py +++ b/api/tests/unit_tests/services/tracing/test_canonical.py @@ -6,7 +6,7 @@ from graphon.entities import WorkflowNodeExecution from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus -from services.tracing.canonical import CanonicalTraceBuilder, root_trace_id +from services.tracing.canonical import CanonicalTraceBuilder, message_trace_id from services.tracing.records import ( MessageTraceInfo, PromptGenerationTraceInfo, @@ -116,7 +116,6 @@ def _workflow_record(*, message_id: str | None, metadata: dict[str, object]) -> file_list=[], invoked_by=None, query="hi" if message_id else "", - parent=None, start_time=STARTED_AT, end_time=STARTED_AT + timedelta(seconds=3), metadata=metadata, @@ -145,7 +144,7 @@ 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 == root_trace_id(None, MESSAGE_ID) == trace_id_for(MESSAGE_ID) + 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 diff --git a/api/tests/unit_tests/services/tracing/test_delivery.py b/api/tests/unit_tests/services/tracing/test_delivery.py index d33f3cce2994ae..3b80818652068d 100644 --- a/api/tests/unit_tests/services/tracing/test_delivery.py +++ b/api/tests/unit_tests/services/tracing/test_delivery.py @@ -6,20 +6,16 @@ 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, TraceParent, WorkflowRunCompleted -from dify_trace_sdk.ids import span_id_for, trace_id_for -from dify_trace_sdk.model import SpanRef, Trace +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 ( - CollectedTrace, ExportJob, ExportRetryQueue, - OuterRunQuery, - OuterRunRoot, StoredTraceDestination, TraceCollector, TraceDeliveryService, @@ -27,7 +23,7 @@ TraceRecordSource, retry_delay_seconds, ) -from services.tracing.errors import SourceNotReadyError, TraceCollectorUnavailableError, TracingStoreUnavailableError +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 @@ -36,9 +32,6 @@ MESSAGE_ID = "6f1c0b52-6ad4-4c4e-9d7c-2a3f0d0c7a11" TENANT_ID = "tenant-1" APP_ID = "app-1" -NESTED_RUN_ID = "nested-run" -TOOL_STARTED_AT = STARTED_AT + timedelta(seconds=2) -SAME_DESTINATION = StoredTraceDestination(provider_id="langfuse", config={"endpoint": "https://collector.example"}) class _Config(ProviderConfig): @@ -99,33 +92,6 @@ def active_destination(self, *, tenant_id: str, app_id: str) -> StoredTraceDesti return self.destinations.get((tenant_id, app_id)) -class _OuterRuns(OuterRunQuery): - """Outer run roots by (tenant, app, run) and node starts by (tenant, app, run, node).""" - - def __init__(self) -> None: - self.roots: dict[tuple[str, str, str], OuterRunRoot] = {} - self.node_starts: dict[tuple[str, str, str, str], datetime] = {} - self.unreadable = False - - @override - def root(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> OuterRunRoot | None: - if self.unreadable: - raise TracingStoreUnavailableError - return self.roots.get((tenant_id, app_id, workflow_run_id)) - - @override - def node_started_at( - self, *, tenant_id: str, app_id: str, workflow_run_id: str, node_execution_id: str - ) -> datetime | None: - return self.node_starts.get((tenant_id, app_id, workflow_run_id, node_execution_id)) - - def add(self, link: TraceParent, root: OuterRunRoot) -> None: - assert link.app_id is not None - assert link.node_execution_id is not None - self.roots[(TENANT_ID, link.app_id, link.workflow_run_id)] = root - self.node_starts[(TENANT_ID, link.app_id, link.workflow_run_id, link.node_execution_id)] = TOOL_STARTED_AT - - class _Retries(ExportRetryQueue): def __init__(self) -> None: self.scheduled: list[tuple[ExportJob, int]] = [] @@ -140,7 +106,6 @@ def __init__(self, *, captures: bool, fails: bool) -> None: self.captures = captures self.fails = fails self.seen: list[tuple[str, BaseTraceInfo]] = [] - self.collected: dict[str, CollectedTrace] = {} @override def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: @@ -150,8 +115,8 @@ def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: return self.captures @override - def collected_trace(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> CollectedTrace | None: - return self.collected.get(workflow_run_id) + def link_managed_parent(self, *, tenant_id: str, info: WorkflowTraceInfo) -> WorkflowTraceInfo: + return info @dataclass @@ -160,13 +125,11 @@ class _Delivery: egress: RecordingEgress destinations: _Destinations retries: _Retries = field(default_factory=_Retries) - outer_runs: _OuterRuns = field(default_factory=_OuterRuns) enterprise: list[BaseTraceInfo] = field(default_factory=list) def service(self, *, enterprise: bool, collector: _Collector | None) -> TraceDeliveryService: return TraceDeliveryService( records=self.records, - outer_runs=self.outer_runs, destinations=self.destinations, providers=InstalledTraceProviders(_Provider), codec=ProviderConfigCodec(WorkspaceBoundCipher()), @@ -189,7 +152,9 @@ def _event() -> TraceEvent: app_id=APP_ID, user_id="account-1", occurred_at=datetime.now(UTC), - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + 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"), ) @@ -230,71 +195,12 @@ def _record(*, workflow_run_id: str | None) -> MessageTraceInfo: ) -def _workflow_event() -> TraceEvent: - return _event().model_copy( - update={"payload": WorkflowRunCompleted(workflow_run_id=NESTED_RUN_ID, conversation_id=None, total_tokens=None)} - ) - - -def _nested_record(parent: TraceParent) -> WorkflowTraceInfo: - return WorkflowTraceInfo( - node_executions=[], - conversation_id=None, - message_started_at=None, - workflow_app_log_id=None, - workflow_id="workflow-1", - tenant_id=TENANT_ID, - workflow_run_id=NESTED_RUN_ID, - workflow_run_elapsed_time=1, - 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="", - parent=parent, - start_time=STARTED_AT + timedelta(seconds=3), - end_time=STARTED_AT + timedelta(seconds=4), - metadata={}, - ) - - -def _link(run_id: str, *, external_trace_id: str | None, parent: TraceParent | None) -> TraceParent: - return TraceParent( - workflow_run_id=run_id, - app_id=f"{run_id}-app", - node_execution_id=f"{run_id}-tool-node", - external_trace_id=external_trace_id, - parent=parent, - ) - - -def _nested_delivery(parent: TraceParent) -> _Delivery: - """A nested run whose outer runs are all persisted and trace where it does.""" - delivery = _delivery(_nested_record(parent), [response(200, {})]) - link: TraceParent | None = parent - while link is not None: - assert link.app_id is not None - delivery.destinations.destinations[(TENANT_ID, link.app_id)] = SAME_DESTINATION - delivery.outer_runs.add(link, OuterRunRoot(root_id=link.workflow_run_id, started_at=STARTED_AT)) - link = link.parent - return delivery - - -def _tool_node(run_id: str) -> SpanRef: - return SpanRef(span_id=span_id_for(f"{run_id}-tool-node"), start_time=TOOL_STARTED_AT) - - 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): SAME_DESTINATION}), + destinations=_Destinations({(TENANT_ID, APP_ID): stored}), ) @@ -459,82 +365,3 @@ def test_destinations_that_cannot_be_used_send_nothing(stored: StoredTraceDestin delivery.service(enterprise=False, collector=None).deliver(_event(), enterprise_delivered=False) assert delivery.sent() == [] - - -def test_nested_run_hangs_from_the_tool_node_in_the_trace_of_its_outer_run() -> None: - delivery = _nested_delivery(_link("outer-run", external_trace_id=None, parent=None)) - - delivery.service(enterprise=False, collector=None).deliver(_workflow_event(), enterprise_delivered=False) - - (trace,) = delivery.sent() - assert trace.trace_id == trace_id_for("outer-run") - assert trace.external_trace_id is None - assert trace.parent == _tool_node("outer-run") - assert trace.trace_root == SpanRef(span_id=span_id_for("outer-run"), start_time=STARTED_AT) - assert trace.root.span_id == span_id_for(NESTED_RUN_ID) - assert trace.root.parent_span_id == span_id_for("outer-run-tool-node") - - -def test_outer_runs_sharing_the_destination_join_the_trace_of_the_outermost_one() -> None: - chatflow = _link("chatflow-run", external_trace_id="caller-trace", parent=None) - delivery = _nested_delivery(_link("outer-run", external_trace_id=None, parent=chatflow)) - message_started_at = STARTED_AT - timedelta(seconds=1) - delivery.outer_runs.roots[(TENANT_ID, "chatflow-run-app", "chatflow-run")] = OuterRunRoot( - root_id="chatflow-message", started_at=message_started_at - ) - - delivery.service(enterprise=False, collector=None).deliver(_workflow_event(), enterprise_delivered=False) - - (trace,) = delivery.sent() - assert (trace.trace_id, trace.external_trace_id) == (trace_id_for("caller-trace"), "caller-trace") - assert trace.parent == _tool_node("outer-run") - assert trace.trace_root == SpanRef(span_id=span_id_for("chatflow-message"), start_time=message_started_at) - - -def test_a_nested_run_starts_its_own_trace_below_an_outer_run_tracing_elsewhere() -> None: - delivery = _nested_delivery(_link("outer-run", external_trace_id=None, parent=None)) - delivery.destinations.destinations[(TENANT_ID, "outer-run-app")] = StoredTraceDestination( - provider_id="langfuse", config={"endpoint": "https://elsewhere.example"} - ) - - delivery.service(enterprise=False, collector=None).deliver(_workflow_event(), enterprise_delivered=False) - - (trace,) = delivery.sent() - assert (trace.trace_id, trace.parent, trace.trace_root) == (trace_id_for(NESTED_RUN_ID), None, None) - - -def test_a_collected_outer_run_lends_its_agent_trace_once_collection_starts() -> None: - collector = _Collector(captures=False, fails=False) - agent_trace = CollectedTrace(trace_id=trace_id_for("agent-trace"), destination=SAME_DESTINATION) - collector.collected["outer-run"] = agent_trace - delivery = _nested_delivery(_link("outer-run", external_trace_id="caller-trace", parent=None)) - - delivery.service(enterprise=False, collector=collector).deliver(_workflow_event(), enterprise_delivered=False) - - (trace,) = delivery.sent() - assert (trace.trace_id, trace.external_trace_id) == (trace_id_for("agent-trace"), None) - assert trace.parent == _tool_node("outer-run") - - -def test_a_collected_outer_run_lends_no_trace_before_collection_starts() -> None: - collector = _Collector(captures=False, fails=False) - collector.collected["outer-run"] = CollectedTrace(trace_id=None, destination=None) - delivery = _nested_delivery(_link("outer-run", external_trace_id=None, parent=None)) - - delivery.service(enterprise=False, collector=collector).deliver(_workflow_event(), enterprise_delivered=False) - - (trace,) = delivery.sent() - assert (trace.trace_id, trace.parent) == (trace_id_for(NESTED_RUN_ID), None) - - -def test_a_nested_run_starts_its_own_trace_while_its_outer_records_are_unavailable() -> None: - unpersisted = _nested_delivery(_link("outer-run", external_trace_id=None, parent=None)) - unpersisted.outer_runs.node_starts.clear() - unreadable = _nested_delivery(_link("outer-run", external_trace_id=None, parent=None)) - unreadable.outer_runs.unreadable = True - - for delivery in (unpersisted, unreadable): - delivery.service(enterprise=False, collector=None).deliver(_workflow_event(), enterprise_delivered=False) - - (trace,) = delivery.sent() - assert (trace.trace_id, trace.parent) == (trace_id_for(NESTED_RUN_ID), None) diff --git a/api/tests/unit_tests/services/tracing/test_recovery.py b/api/tests/unit_tests/services/tracing/test_recovery.py index f8db25d6df9163..40c5225030c0dd 100644 --- a/api/tests/unit_tests/services/tracing/test_recovery.py +++ b/api/tests/unit_tests/services/tracing/test_recovery.py @@ -272,7 +272,9 @@ def trace_event(run: TraceRun, payload: TracePayload) -> TraceEvent: app_id=run.app_id, user_id=None, occurred_at=datetime.now(UTC), - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), payload=payload, ) @@ -550,7 +552,9 @@ def send(message: Message, *, application_generate_entity: ChatAppGenerateEntity 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=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ), ) stop = MessageQueueMessage( 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 3104feba413be4..ccc648d5d37fee 100644 --- a/api/tests/unit_tests/services/tracing/test_runtime_bridge.py +++ b/api/tests/unit_tests/services/tracing/test_runtime_bridge.py @@ -18,11 +18,10 @@ ModerationChecked, TraceCorrelation, TraceEvent, - TraceParent, TracePayload, WorkflowRunCompleted, ) -from dify_trace_sdk.ids import span_id_for, span_uuid, trace_id_for, trace_uuid +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 @@ -311,7 +310,9 @@ def trace_event(run: TraceRun, payload: TracePayload) -> TraceEvent: app_id=run.app_id, user_id=None, occurred_at=datetime.now(UTC), - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), payload=payload, ) @@ -348,7 +349,6 @@ def terminal(collection: Collection) -> None: session_id=run.session_id, user_id=None, parent=None, - trace_root=None, spans=(root,), ) capture(collection.repository, run.id, trace) @@ -384,10 +384,8 @@ def setup(sessions: sessionmaker[Session], name: str, mode: AppMode, *, snapshot 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 - records = record_source(sessions) TraceDeliveryService( - records=records, - outer_runs=records, + records=record_source(sessions), destinations=SQLAlchemyAppTracingConfigRepository(session_factory=sessions), providers=InstalledTraceProviders(LangfuseProvider), codec=ProviderConfigCodec(WorkspaceBoundCipher()), @@ -422,15 +420,8 @@ def workflow_event(row: WorkflowRun, parent: WorkflowRun | None) -> TraceEvent: correlation=TraceCorrelation( external_trace_id=None, session_id=None, - parent=TraceParent( - workflow_run_id=parent.id, - app_id=parent.app_id, - node_execution_id=identifier("parent-tool-node"), - external_trace_id=None, - parent=None, - ) - if parent - else 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), ) @@ -486,46 +477,6 @@ def workflow(case: Case, name: str, *, has_agent: bool, snapshot: bool) -> tuple return row, trace_run -def tool_node(case: Case, row: WorkflowRun) -> None: - """Persist the tool node execution of ``row`` that starts the nested runs of these cases.""" - with case.collection.sessions.begin() as session: - node = WorkflowNodeExecutionModel( - id=identifier("parent-tool-node"), - node_execution_id=identifier("parent-tool-node"), - tenant_id=row.tenant_id, - app_id=row.app_id, - workflow_id=row.workflow_id, - workflow_run_id=row.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=row.created_by, - created_by_role=CreatorUserRole.ACCOUNT, - ) - node.created_at = row.created_at + timedelta(seconds=1) - session.add(node) - - -def nested_root(destination: RecordingEgress, trace_id: str, child: WorkflowRun) -> dict[str, object]: - """The Langfuse observation of a nested run's root, after checking it was sent as a fragment.""" - events = ingested(destination) - assert all(item["type"] != "trace-create" for item in events) - bodies = [item["body"] for item in events if isinstance(item["body"], dict)] - (root,) = [body for body in bodies if body["id"] == str(span_uuid(trace_id, span_id_for(child.id)))] - assert root["traceId"] == str(trace_uuid(trace_id)) - assert root["parentObservationId"] == str(span_uuid(trace_id, span_id_for(identifier("parent-tool-node")))) - return root - - 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} @@ -554,7 +505,6 @@ def workflow_info(row: WorkflowRun, parent: WorkflowRun | None) -> WorkflowTrace file_list=[], invoked_by=None, query="", - parent=None, metadata=metadata, start_time=row.created_at, end_time=row.finished_at, @@ -678,12 +628,11 @@ def test_auxiliary_events_resolve_owner_without_bypassing_content_policy( assert run_count(case.collection.sessions) == 0 @staticmethod - def test_plain_nested_workflow_hangs_from_the_tool_node_of_its_parent( + def test_plain_nested_workflow_keeps_its_parent_context_for_the_provider( sqlite_session_factory: sessionmaker[Session], ) -> None: parent_case = setup(sqlite_session_factory, "plain-parent", AppMode.WORKFLOW, snapshot=False) parent, _ = workflow(parent_case, "plain-parent", has_agent=False, snapshot=False) - tool_node(parent_case, parent) child_case = setup(sqlite_session_factory, "plain-child", AppMode.WORKFLOW, snapshot=False) child, _ = workflow(child_case, "plain-child", has_agent=False, snapshot=False) add_app_owner(child_case.collection, "plain-child") @@ -691,19 +640,15 @@ def test_plain_nested_workflow_hangs_from_the_tool_node_of_its_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.collected_trace( - tenant_id=parent.tenant_id, app_id=parent.app_id, workflow_run_id=parent.id - ) - is None - ) + assert child_case.service.link_managed_parent(tenant_id=child.tenant_id, info=record) is record destination = exported() deliver(child_case, event, destination) - root = nested_root(destination, trace_id_for(parent.id), child) - metadata = root["metadata"] - assert isinstance(metadata, dict) - assert metadata["parent_workflow_run_id"] == parent.id - assert metadata["parent_node_execution_id"] == identifier("parent-tool-node") + 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 @@ -729,40 +674,32 @@ def test_plain_chatflow_auxiliary_events_keep_the_provider_route( assert len(destination.requests) == 1 @staticmethod - def test_provider_exported_child_of_a_collected_parent_hangs_in_its_agent_trace( + def test_provider_exported_child_of_managed_parent_names_the_collected_parent( sqlite_session_factory: sessionmaker[Session], ) -> None: 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 - tool_node(parent_case, parent) child_case = setup(sqlite_session_factory, "provider-child", AppMode.WORKFLOW, snapshot=False) child, _ = workflow(child_case, "provider-child", has_agent=False, snapshot=False) - collected = child_case.service.collected_trace( - tenant_id=parent.tenant_id, app_id=parent.app_id, workflow_run_id=parent.id - ) - assert collected is not None - assert collected.trace_id == parent_trace.trace_id - assert collected.destination is not None - assert collected.destination.provider_id == parent_trace.provider + info = workflow_info(child, parent) + 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) - nested_root(destination, parent_trace.trace_id, child) - - @staticmethod - def test_a_collected_parent_lends_no_trace_before_its_collection_starts( - sqlite_session_factory: sessionmaker[Session], - ) -> None: - parent_case = setup(sqlite_session_factory, "pending-parent", AppMode.WORKFLOW, snapshot=False) - parent, _ = workflow(parent_case, "pending-parent", has_agent=True, snapshot=False) - - collected = parent_case.service.collected_trace( - tenant_id=parent.tenant_id, app_id=parent.app_id, workflow_run_id=parent.id - ) - - assert collected is not None - assert (collected.trace_id, collected.destination) == (None, None) + 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( @@ -1477,7 +1414,6 @@ def test_hidden_child_skip_keeps_normal_root_agent_trace_assemblable( file_list=[], invoked_by=None, query="", - parent=None, metadata={"app_id": ids["root"]}, start_time=workflow.created_at, end_time=workflow.finished_at, diff --git a/api/tests/unit_tests/services/tracing/test_scopes.py b/api/tests/unit_tests/services/tracing/test_scopes.py index ffd9c410991725..f269375ea3c8b0 100644 --- a/api/tests/unit_tests/services/tracing/test_scopes.py +++ b/api/tests/unit_tests/services/tracing/test_scopes.py @@ -6,7 +6,9 @@ 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=None) +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") diff --git a/api/tests/unit_tests/services/tracing/test_service.py b/api/tests/unit_tests/services/tracing/test_service.py index 5a07dd21472964..1dfe6b4a4a3276 100644 --- a/api/tests/unit_tests/services/tracing/test_service.py +++ b/api/tests/unit_tests/services/tracing/test_service.py @@ -354,7 +354,9 @@ def trace_event(case: _ApplicationTraceCollectionCollection, payload: TracePaylo app_id=case.run.app_id, user_id=None, occurred_at=datetime.now(UTC), - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), payload=payload, ) @@ -425,7 +427,6 @@ def _application_trace_collection_terminal(collection: _ApplicationTraceCollecti session_id=run.session_id, user_id=None, parent=None, - trace_root=None, spans=(root,), ) with collection.change() as current: 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 86a22362474a6b..d02af7d055d76e 100644 --- a/api/tests/unit_tests/tasks/test_application_trace_task.py +++ b/api/tests/unit_tests/tasks/test_application_trace_task.py @@ -128,7 +128,9 @@ def test_worker_records_carry_no_enterprise_names(sqlite_session_factory: sessio app_id=case.run.app_id, user_id=None, occurred_at=datetime.now(UTC), - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + 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), ) ) diff --git a/api/tests/unit_tests/tasks/test_trace_delivery_task.py b/api/tests/unit_tests/tasks/test_trace_delivery_task.py index afb4ff04cd0c3a..9e70a12ba7d04d 100644 --- a/api/tests/unit_tests/tasks/test_trace_delivery_task.py +++ b/api/tests/unit_tests/tasks/test_trace_delivery_task.py @@ -15,11 +15,8 @@ from services.tracing.adapters import payload_key from services.tracing.canonical import CanonicalTraceBuilder from services.tracing.delivery import ( - CollectedTrace, ExportJob, ExportRetryQueue, - OuterRunQuery, - OuterRunRoot, StoredTraceDestination, TraceCollector, TraceDeliveryService, @@ -28,7 +25,7 @@ ) from services.tracing.errors import SourceNotFoundError, SourceNotReadyError from services.tracing.providers import ProviderConfigCodec -from services.tracing.records import BaseTraceInfo, GenerateNameTraceInfo +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 @@ -91,20 +88,8 @@ def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: raise RuntimeError("ledger unavailable") @override - def collected_trace(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> CollectedTrace | None: - raise AssertionError("no nested workflow run is delivered") - - -class _NoOuterRuns(OuterRunQuery): - @override - def root(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> OuterRunRoot | None: - raise AssertionError("no nested workflow run is delivered") - - @override - def node_started_at( - self, *, tenant_id: str, app_id: str, workflow_run_id: str, node_execution_id: str - ) -> datetime | None: - raise AssertionError("no nested workflow run is delivered") + def link_managed_parent(self, *, tenant_id: str, info: WorkflowTraceInfo) -> WorkflowTraceInfo: + return info class _Delivery: @@ -115,7 +100,6 @@ def __init__(self) -> None: def service(self, collector: TraceCollector | None) -> TraceDeliveryService: return TraceDeliveryService( records=self.records, - outer_runs=_NoOuterRuns(), destinations=self.destinations, providers=InstalledTraceProviders(), codec=ProviderConfigCodec(WorkspaceBoundCipher()), @@ -144,7 +128,9 @@ def _event() -> TraceEvent: app_id="app-1", user_id=None, occurred_at=datetime.now(UTC), - correlation=TraceCorrelation(external_trace_id=None, session_id=None, parent=None), + 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), ) From 644f0c8b1494fa3cb1e812040fd125a31c860a8d Mon Sep 17 00:00:00 2001 From: -LAN- Date: Thu, 29 Oct 2026 19:31:32 +0800 Subject: [PATCH 07/17] fix(tracing): include nested workflow tool frames in exports Load persisted workflow-tool descendants through their admitted caller, tenant, and root run. Scope predecessor and container lookup to each invocation so repeated calls and retries retain the correct parent. Inject workflow readers from the tracing composition root, hydrate SQL and LogStore payloads through their repository contracts, and scope offloaded files to the execution tenant before reading storage outside the database session. --- .../application_services/tracing.py | 17 +- ..._api_workflow_node_execution_repository.py | 8 + .../api_workflow_node_execution_repository.py | 8 + ..._api_workflow_node_execution_repository.py | 36 +++- api/repositories/trace_record_repository.py | 108 +++++++--- api/services/tracing/canonical.py | 2 +- api/services/tracing/delivery.py | 3 + api/services/tracing/hierarchy.py | 48 +++-- api/services/tracing/records.py | 4 +- .../inner_api/agent/test_tracing.py | 9 +- .../enterprise/telemetry/test_draft_trace.py | 9 +- .../test_application_trace_repository.py | 17 +- ..._api_workflow_node_execution_repository.py | 66 ++++++ .../test_trace_record_repository.py | 204 +++++++++++++++++- .../services/tracing/test_canonical.py | 98 +++++++++ .../services/tracing/test_recovery.py | 9 +- .../services/tracing/test_runtime_bridge.py | 9 +- .../services/tracing/test_service.py | 9 +- .../tasks/test_application_trace_task.py | 9 +- 19 files changed, 616 insertions(+), 57 deletions(-) diff --git a/api/extensions/application_services/tracing.py b/api/extensions/application_services/tracing.py index 26e8b3d61ecc9a..b4ef5cd0e4189e 100644 --- a/api/extensions/application_services/tracing.py +++ b/api/extensions/application_services/tracing.py @@ -8,6 +8,7 @@ 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, @@ -46,6 +47,8 @@ def build_tracing_services(*, database_client: sessionmaker[Session]) -> Tracing 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, @@ -61,7 +64,13 @@ def build_tracing_services(*, database_client: sessionmaker[Session]) -> Tracing run_query=agent_trace_runs, sources=agent_trace_runs, # App, workspace, and credential names are enterprise telemetry dimensions only. - records=TraceRecordRepository(database_client, files_url=dify_config.FILES_URL, enterprise_details=False), + 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, @@ -75,7 +84,11 @@ def build_tracing_services(*, database_client: sessionmaker[Session]) -> Tracing ), delivery=TraceDeliveryService( records=TraceRecordRepository( - database_client, files_url=dify_config.FILES_URL, enterprise_details=enterprise_telemetry + database_client, + workflow_runs=workflow_runs, + node_executions=node_executions, + files_url=dify_config.FILES_URL, + enterprise_details=enterprise_telemetry, ), destinations=configs, providers=providers, 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..24e8b125464967 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 @@ -161,6 +161,14 @@ 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.""" diff --git a/api/repositories/api_workflow_node_execution_repository.py b/api/repositories/api_workflow_node_execution_repository.py index 540cf689109797..20be54e9d522ad 100644 --- a/api/repositories/api_workflow_node_execution_repository.py +++ b/api/repositories/api_workflow_node_execution_repository.py @@ -184,6 +184,14 @@ 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 get_workflow_tool_executions( self, tenant_id: str, diff --git a/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py b/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py index 219dbd32721cb6..a95afaf258d730 100644 --- a/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py @@ -5,6 +5,7 @@ using SQLAlchemy 2.0 style queries for WorkflowNodeExecutionModel operations. """ +import json from collections.abc import Mapping, Sequence from datetime import datetime from typing import Any, cast, override @@ -15,7 +16,13 @@ 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, WorkflowNodeExecutionSnapshot, @@ -52,11 +59,36 @@ 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 with self._session_maker() as session: - return execution.load_full_process_data(session, storage) + key = session.scalar( + select(UploadFile.key).where( + UploadFile.id == offload.file_id, UploadFile.tenant_id == execution.tenant_id + ) + ) + assert key is not None, f"UploadFile with id {offload.file_id} should exist in the execution's tenant" + return json.loads(storage.load(key)) @override def get_node_last_execution( diff --git a/api/repositories/trace_record_repository.py b/api/repositories/trace_record_repository.py index 0b6bd57a1f3522..cc4a7c8f07f4f8 100644 --- a/api/repositories/trace_record_repository.py +++ b/api/repositories/trace_record_repository.py @@ -2,11 +2,13 @@ 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. Rows -become plain records before their session closes. +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 @@ -26,16 +28,21 @@ from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker -from core.repositories import DifyCoreRepositoryFactory from graphon.entities import WorkflowNodeExecution -from graphon.enums import WorkflowExecutionStatus -from models import Account, App, Tenant +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, WorkflowNodeExecutionTriggeredFrom, WorkflowRun -from repositories.factory import DifyAPIRepositoryFactory +from models.workflow import WorkflowAppLog, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun +from repositories.api_workflow_node_execution_repository import DifyAPIWorkflowNodeExecutionRepository +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 @@ -124,8 +131,18 @@ def _streaming_metrics(message: Message) -> dict[str, Any]: class TraceRecordRepository(TraceRecordSource): - def __init__(self, session_factory: sessionmaker[Session], *, files_url: str, enterprise_details: bool) -> None: + 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 @@ -278,9 +295,8 @@ def _message(self, event: TraceEvent, payload: MessageCompleted) -> MessageTrace def _workflow(self, event: TraceEvent, payload: WorkflowRunCompleted) -> WorkflowTraceInfo: if event.app_id is None: raise SourceNotFoundError - runs = DifyAPIRepositoryFactory.create_api_workflow_run_repository(self._sessions) with tracing_store_errors(): - run = runs.get_workflow_run_by_id( + 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). @@ -307,7 +323,7 @@ def _workflow(self, event: TraceEvent, payload: WorkflowRunCompleted) -> Workflo else None ) app_name, workspace_name = self._names(session, event.app_id, event.tenant_id) - executions = self._node_executions(run) + 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 [] @@ -338,6 +354,7 @@ def _workflow(self, event: TraceEvent, payload: WorkflowRunCompleted) -> Workflo 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, @@ -362,20 +379,63 @@ def _workflow(self, event: TraceEvent, payload: WorkflowRunCompleted) -> Workflo invoked_by=_user_of(metadata), ) - def _node_executions(self, run: WorkflowRun) -> Sequence[WorkflowNodeExecution]: - # Repositories take an identity only to attribute writes. This read needs no member of the - # workspace, so a run of an app whose creator left it is still traced. - reader = Account(name="Trace assembly", email="trace-assembly@localhost") - reader.id = run.created_by - repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( - session_factory=self._sessions, - tenant_id=run.tenant_id, - user=reader, - app_id=run.app_id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) + def _node_executions(self, run: WorkflowRun) -> tuple[Sequence[WorkflowNodeExecution], Mapping[str, str]]: + repository = self._node_executions_repository + executions: list[WorkflowNodeExecution] = [] + tool_parents: dict[str, str] = {} + seen: set[str] = set() with tracing_store_errors(): - return repository.get_by_workflow_execution(workflow_execution_id=run.id) + 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) + executions.append(self._node_execution(row, repository)) + 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 + ) + ) + return executions, tool_parents + + @staticmethod + def _node_execution( + row: WorkflowNodeExecutionModel, repository: DifyAPIWorkflowNodeExecutionRepository + ) -> 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=repository.load_full_inputs(row), + process_data=repository.load_full_process_data(row), + outputs=repository.load_full_outputs(row), + 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]: diff --git a/api/services/tracing/canonical.py b/api/services/tracing/canonical.py index 505fb01a7a645b..0aa1d484b9d68a 100644 --- a/api/services/tracing/canonical.py +++ b/api/services/tracing/canonical.py @@ -221,7 +221,7 @@ def _message(self, record: MessageTraceInfo) -> Trace: def _workflow(self, record: WorkflowTraceInfo) -> Trace: executions = record.node_executions - hierarchy = build_workflow_hierarchy(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) diff --git a/api/services/tracing/delivery.py b/api/services/tracing/delivery.py index 9e91e3e0a80d22..8727861a3093d0 100644 --- a/api/services/tracing/delivery.py +++ b/api/services/tracing/delivery.py @@ -43,6 +43,9 @@ 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. diff --git a/api/services/tracing/hierarchy.py b/api/services/tracing/hierarchy.py index a42885b965148a..77f0af85a272ea 100644 --- a/api/services/tracing/hierarchy.py +++ b/api/services/tracing/hierarchy.py @@ -1,9 +1,8 @@ """Deterministically reconstruct the span hierarchy of one workflow run. -Parents come from graph predecessors and from iteration/loop containers; +Parents come from graph predecessors, workflow tools, and iteration/loop containers; each iteration or loop round gets a synthetic wrapper. Ambiguous repeated node -IDs and cyclic predecessor data are dropped rather than guessed, so those spans -fall back to the workflow root. +IDs within an invocation and cyclic predecessor data are dropped rather than guessed. """ import re @@ -55,18 +54,29 @@ def finished_at(execution: WorkflowNodeExecution) -> datetime: return execution.finished_at or execution.created_at + timedelta(seconds=execution.elapsed_time or 0) -def _unique_execution_by_node_id(executions: Sequence[WorkflowNodeExecution]) -> dict[str, str]: - result: dict[str, str] = {} - ambiguous: set[str] = set() +type _NodeKey = tuple[str, str | None, str | None, str] + + +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[WorkflowNodeExecution], tool_parents: Mapping[str, str] +) -> dict[_NodeKey, str]: + result: dict[_NodeKey, str] = {} + ambiguous: set[_NodeKey] = set() for item in executions: - if item.node_id in ambiguous: + key = _node_key(item, item.node_id, tool_parents) + if key in ambiguous: continue - previous = result.get(item.node_id) + previous = result.get(key) if previous is None: - result[item.node_id] = item.id + result[key] = item.id elif previous != item.id: - result.pop(item.node_id, None) - ambiguous.add(item.node_id) + result.pop(key, None) + ambiguous.add(key) return result @@ -97,22 +107,28 @@ def _remove_cycles(parent_by_execution_id: dict[str, str]) -> None: parent_by_execution_id.pop(execution, None) -def build_workflow_hierarchy(executions: Sequence[WorkflowNodeExecution]) -> WorkflowHierarchy: - 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: parent_execution_id = ( - execution_by_node_id.get(item.predecessor_node_id) if item.predecessor_node_id is not None 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.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 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 @@ -126,7 +142,7 @@ def build_workflow_hierarchy(executions: Sequence[WorkflowNodeExecution]) -> Wor 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) + 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) diff --git a/api/services/tracing/records.py b/api/services/tracing/records.py index 0e1a911a916304..4737be9a441e0d 100644 --- a/api/services/tracing/records.py +++ b/api/services/tracing/records.py @@ -11,7 +11,7 @@ from enum import StrEnum from typing import Any -from pydantic import BaseModel, ConfigDict, StrictStr, ValidationError +from pydantic import BaseModel, ConfigDict, Field, StrictStr, ValidationError from graphon.entities import WorkflowNodeExecution @@ -95,6 +95,8 @@ class WorkflowTraceInfo(BaseTraceInfo): 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.""" 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 2def94dddb1b3f..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 @@ -13,6 +13,7 @@ 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 @@ -53,7 +54,13 @@ def bind_collector( runs=collection.repository, run_query=collection.repository, sources=collection.repository, - records=TraceRecordRepository(collection.sessions, files_url="https://files.test", enterprise_details=False), + 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()), 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 90502db2cca719..abbf4fe675fe63 100644 --- a/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py +++ b/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py @@ -24,6 +24,7 @@ WorkflowNodeExecutionStatus, WorkflowNodeExecutionTriggeredFrom, ) +from repositories.factory import DifyAPIRepositoryFactory from repositories.trace_record_repository import TraceRecordRepository from services.tracing.records import DraftNodeExecutionTrace @@ -527,7 +528,13 @@ def test_node_data_feeds_into_draft_node_execution_trace( (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, files_url="https://files.example.com", enterprise_details=False + 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 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 2cc696c336ba2a..068cd3d95fd6fa 100644 --- a/api/tests/unit_tests/repositories/test_application_trace_repository.py +++ b/api/tests/unit_tests/repositories/test_application_trace_repository.py @@ -31,6 +31,7 @@ from models.model import App, AppMode, Message, TraceAppConfig from models.workflow import WorkflowNodeExecutionModel, WorkflowRun 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 @@ -76,7 +77,13 @@ def identifier(value: str) -> str: def trace_records(sessions: sessionmaker[Session]) -> TraceRecordRepository: - return TraceRecordRepository(sessions, files_url="https://files.test", enterprise_details=False) + 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( @@ -479,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" 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..7d0829a38ee63e 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 @@ -10,8 +10,11 @@ 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 +29,69 @@ from tests.helpers.unit.config_override import apply_config_overrides +@pytest.mark.parametrize("foreign_tenant", [False, True]) +@pytest.mark.parametrize("payload_type", list(ExecutionOffLoadType)) +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, +) -> 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( + 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] + assert load(execution) == {"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"): + load(execution) + else: + assert load(execution) == {"complete": "payload"} + finally: + event.remove(sqlite_session_factory, "after_begin", track_reader) + + @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 index 51f882d3fdaa62..d1e9b6f1cd0b0d 100644 --- a/api/tests/unit_tests/repositories/test_trace_record_repository.py +++ b/api/tests/unit_tests/repositories/test_trace_record_repository.py @@ -1,3 +1,4 @@ +import json from datetime import UTC, datetime, timedelta from decimal import Decimal @@ -9,20 +10,47 @@ TraceCorrelation, TraceEvent, TracePayload, + WorkflowRunCompleted, ) -from sqlalchemy import create_engine +from sqlalchemy import create_engine, text from sqlalchemy.orm import Session, sessionmaker -from models.enums import ConversationFromSource, ConversationStatus, EndUserType, InvokeFrom +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 +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 @@ -91,8 +119,77 @@ def records(sqlite_session: Session) -> None: 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, files_url="https://files.example.com", enterprise_details=False) + 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: @@ -235,3 +332,102 @@ def test_an_unreachable_database_is_reported_as_unavailable() -> None: 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/services/tracing/test_canonical.py b/api/tests/unit_tests/services/tracing/test_canonical.py index 1c22b6f305b150..c4d9633dc73ad3 100644 --- a/api/tests/unit_tests/services/tracing/test_canonical.py +++ b/api/tests/unit_tests/services/tracing/test_canonical.py @@ -1,6 +1,7 @@ from datetime import UTC, datetime, timedelta from decimal import Decimal +import pytest from dify_trace_sdk.ids import span_id_for, trace_id_for from dify_trace_sdk.model import SpanKind, SpanStatus @@ -224,6 +225,103 @@ def test_failed_node_marks_its_span_as_error() -> None: 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")) diff --git a/api/tests/unit_tests/services/tracing/test_recovery.py b/api/tests/unit_tests/services/tracing/test_recovery.py index 40c5225030c0dd..89fbcbd3e9a714 100644 --- a/api/tests/unit_tests/services/tracing/test_recovery.py +++ b/api/tests/unit_tests/services/tracing/test_recovery.py @@ -63,6 +63,7 @@ from models.workflow import WorkflowRun, WorkflowType 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 @@ -146,7 +147,13 @@ def send(self, config: LangfuseConfig, trace: Trace) -> SendResult: def record_source(sessions: sessionmaker[Session]) -> TraceRecordRepository: - return TraceRecordRepository(sessions, files_url="https://files.test", enterprise_details=False) + 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( 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 ccc648d5d37fee..b3c8c810c71369 100644 --- a/api/tests/unit_tests/services/tracing/test_runtime_bridge.py +++ b/api/tests/unit_tests/services/tracing/test_runtime_bridge.py @@ -61,6 +61,7 @@ ) 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 @@ -138,7 +139,13 @@ def app_tracing(self, *, tenant_id: str, app_id: str) -> AppTracingSnapshot | No def record_source(sessions: sessionmaker[Session]) -> TraceRecordRepository: - return TraceRecordRepository(sessions, files_url="https://files.test", enterprise_details=False) + 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( diff --git a/api/tests/unit_tests/services/tracing/test_service.py b/api/tests/unit_tests/services/tracing/test_service.py index 1dfe6b4a4a3276..d6e26367c291c8 100644 --- a/api/tests/unit_tests/services/tracing/test_service.py +++ b/api/tests/unit_tests/services/tracing/test_service.py @@ -41,6 +41,7 @@ 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 @@ -195,7 +196,13 @@ def decrypt(self, tenant_id: str, tokens: Sequence[str]) -> NoReturn: def trace_records(sessions: sessionmaker[Session]) -> TraceRecordRepository: - return TraceRecordRepository(sessions, files_url="https://files.test", enterprise_details=False) + 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 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 d02af7d055d76e..da70d0865f7961 100644 --- a/api/tests/unit_tests/tasks/test_application_trace_task.py +++ b/api/tests/unit_tests/tasks/test_application_trace_task.py @@ -18,6 +18,7 @@ from models.application_trace import ApplicationTraceRun 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 @@ -52,7 +53,13 @@ def bind_worker( runs=collection.repository, run_query=collection.repository, sources=collection.repository, - records=TraceRecordRepository(collection.sessions, files_url="https://files.test", enterprise_details=False), + 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()), From 957bc4e6aaebbea3123e2112affd902059293d02 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Thu, 29 Oct 2026 19:39:07 +0800 Subject: [PATCH 08/17] test(tracing): exercise persisted offloads in integration coverage Replace the removed model-helper stub with a committed upload file owned by the execution tenant. Keep the storage transport isolated while checking that the repository resolves and hydrates the real offload reference. --- ..._api_workflow_node_execution_repository.py | 47 +++++++++++++++---- 1 file changed, 38 insertions(+), 9 deletions(-) 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)) From 882b93e9fb4dc65dfd2963559e1194cd22963d2f Mon Sep 17 00:00:00 2001 From: -LAN- Date: Fri, 30 Oct 2026 07:05:01 +0800 Subject: [PATCH 09/17] fix(tracing): preserve existing Langfuse configurations Retain verified project URLs from legacy configurations so project, trace, and observation links work without saving credentials again. Restore Langfuse chat-message formatting for prompts and responses while preserving structured workflow data and multimodal content. Cover legacy encrypted configuration reads, client-supplied derived-field rejection, link generation, and prompt export compatibility. --- .../src/dify_trace_langfuse/config.py | 7 ++ .../src/dify_trace_langfuse/provider.py | 33 +++++++-- .../src/dify_trace_langfuse/test_config.py | 8 ++- .../src/dify_trace_langfuse/test_provider.py | 72 +++++++++++++++++++ .../services/tracing/test_provider_configs.py | 23 ++++++ 5 files changed, 135 insertions(+), 8 deletions(-) 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 ff05d84411f9be..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 @@ -9,8 +9,15 @@ class LangfuseConfig(ProviderConfig): 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(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/provider.py b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/provider.py index b587bba7c8d837..620ecd12388c87 100644 --- a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/provider.py +++ b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/provider.py @@ -23,6 +23,7 @@ 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 @@ -35,6 +36,20 @@ 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 @@ -73,9 +88,9 @@ def verify(self, config: LangfuseConfig) -> LangfuseConfig: return config.model_copy(update={"project_key": project_id}) @override - def project_url(self, config: LangfuseConfig) -> str | None: + def project_url(self, config: LangfuseConfig) -> str: host = config.host.rstrip("/") - return f"{host}/project/{config.project_key}" if config.project_key else host + 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: @@ -101,7 +116,11 @@ def send(self, config: LangfuseConfig, trace: Trace) -> SendResult: # 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)}/traces/{quote(trace_key, safe='')}" if config.project_key else None + 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)}" @@ -130,8 +149,8 @@ def _events(self, trace: Trace, trace_key: str) -> list[dict[str, object]]: "timestamp": _iso(root.start_time), "userId": trace.user_id, "sessionId": trace.session_id, - "input": root.inputs, - "output": root.outputs, + "input": _messages(root.inputs, role="user"), + "output": _messages(root.outputs, role="assistant"), "metadata": {**root.attributes, "dify_trace_id": trace.trace_id}, }, ) @@ -155,8 +174,8 @@ def _events(self, trace: Trace, trace_key: str) -> list[dict[str, object]]: "name": span.name, "startTime": _iso(span.start_time), "endTime": _iso(span.end_time), - "input": span.inputs, - "output": span.outputs, + "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, 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 d836f678671f14..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 @@ -5,7 +5,13 @@ def test_config_secrets_and_defaults() -> None: assert LangfuseConfig.secret_fields() == {"public_key", "secret_key"} - assert LangfuseConfig.derived_fields() == {"project_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_provider.py b/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_provider.py index 0c04bf5edb88cc..8227f931cca99e 100644 --- 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 @@ -12,6 +12,7 @@ ) 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") @@ -39,6 +40,30 @@ def test_verify_resolves_the_project() -> None: 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"), [ @@ -73,10 +98,57 @@ def test_message_trace_becomes_trace_and_generation() -> None: 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() diff --git a/api/tests/unit_tests/services/tracing/test_provider_configs.py b/api/tests/unit_tests/services/tracing/test_provider_configs.py index b1922e92483a73..95eea6f1329bf7 100644 --- a/api/tests/unit_tests/services/tracing/test_provider_configs.py +++ b/api/tests/unit_tests/services/tracing/test_provider_configs.py @@ -1,5 +1,6 @@ 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 @@ -55,6 +56,26 @@ def test_presentation_never_returns_secrets_or_derived_fields() -> None: 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) @@ -67,10 +88,12 @@ def test_new_configs_are_verified_and_store_derived_fields() -> None: "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" From 4af0d787fd0c924426d28b74f8c2112a4e2bb3d4 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Fri, 30 Oct 2026 14:29:41 +0800 Subject: [PATCH 10/17] fix(tracing): reconcile migration and resume coverage after rebase Place the tracing configuration migration after the current main schema head so upgrades retain a single migration path. Preserve the end-user identity in the resumed workflow fixture when replaying the earlier merge resolutions. --- ..._1200-53c10126d551_unique_trace_app_config_provider.py | 4 ++-- .../core/app/apps/advanced_chat/test_app_generator.py | 8 +++++++- 2 files changed, 9 insertions(+), 3 deletions(-) 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 index e81783b3329ff1..78cd9b7c8efe14 100644 --- 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 @@ -1,7 +1,7 @@ """keep one tracing configuration per app and provider Revision ID: 53c10126d551 -Revises: b9d3f7a5c102 +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 @@ -13,7 +13,7 @@ from alembic import op revision = "53c10126d551" -down_revision = "b9d3f7a5c102" +down_revision = "0cc3158b50b0" branch_labels = None depends_on = None 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 9f2f0bafc60747..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 @@ -33,7 +33,12 @@ ) 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, WorkflowAppGenerateEntity +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 @@ -1967,6 +1972,7 @@ def _make_workflow_generate_entity() -> WorkflowAppGenerateEntity: inputs={}, files=[], user_id="user", + user_from=UserFrom.END_USER, stream=False, invoke_from=InvokeFrom.SERVICE_API, workflow_execution_id="paused-run", From 208b6709f1a8ee57d5a9ef93a2b2a43886408c85 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Sat, 31 Oct 2026 01:29:09 +0800 Subject: [PATCH 11/17] fix(tracing): isolate task queues during rolling upgrades Route tracing delivery, retries, Agent exports, and sandbox metering through OPS_TRACE_QUEUE, defaulting to ops_trace_v2. Align Docker and development worker subscriptions so old workers can finish legacy tasks without consuming new payloads. Preserve legacy stored payloads while old workers drain. Document producer shutdown, pending task checks, and the optional return to ops_trace with new workers consuming both queues until retries finish. --- .devcontainer/post_create_command.sh | 2 +- api/.env.example | 2 + api/configs/feature/__init__.py | 5 ++ api/docker/entrypoint.sh | 4 +- api/extensions/ext_celery.py | 7 +- api/providers/trace/README.md | 36 +++++++- api/schedule/collect_agent_sandbox_usage.py | 1 - api/services/tracing/adapters.py | 7 +- api/tasks/application_trace_task.py | 6 +- api/tasks/ops_trace_task.py | 3 +- api/tasks/trace_delivery_task.py | 29 +++---- .../test_knowledge_fs_worker_deployment.py | 25 ++---- api/tests/helpers/unit/worker_entrypoint.py | 20 +++++ .../unit_tests/configs/test_app_config.py | 7 ++ .../unit_tests/extensions/test_ext_celery.py | 87 ++++++++++++++++++- .../test_collect_agent_sandbox_usage.py | 4 +- .../services/tracing/test_adapters.py | 11 ++- .../tasks/test_trace_delivery_task.py | 6 +- .../unit_tests/test_worker_entrypoint.py | 31 ++++--- dev/start-worker | 6 +- docker/README.md | 12 +++ docker/envs/core-services/shared.env.example | 4 +- 22 files changed, 241 insertions(+), 74 deletions(-) create mode 100644 api/tests/helpers/unit/worker_entrypoint.py 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/api/.env.example b/api/.env.example index 4ed6e896513124..21ef0068c9ac62 100644 --- a/api/.env.example +++ b/api/.env.example @@ -118,6 +118,8 @@ CELERY_BROKER_URL=redis://:difyai123456@localhost:${REDIS_PORT}/1 CELERY_BACKEND=redis # Ops trace configuration +# 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 diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 67d4168d38f99a..1371027fd80db2 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -1420,6 +1420,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." ) 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/extensions/ext_celery.py b/api/extensions/ext_celery.py index faccd9e5782e55..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": @@ -210,7 +215,7 @@ def __call__(self, *args: object, **kwargs: object) -> object: "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", # drains trace tasks enqueued before the pipeline upgrade + "tasks.ops_trace_task", # retired endpoint; old workers must drain pre-upgrade tasks ] day = dify_config.CELERY_BEAT_SCHEDULER_TIME diff --git a/api/providers/trace/README.md b/api/providers/trace/README.md index 35749e6b6e2519..984c113b3b09da 100644 --- a/api/providers/trace/README.md +++ b/api/providers/trace/README.md @@ -84,7 +84,7 @@ Stored configurations keep the provider's fields at the top level, with `secret( - moderation, tools, retrieval, conversation naming and suggested questions. A scope exists only when the app has tracing enabled or enterprise telemetry is on. -2. `tasks.trace_delivery_task` hands each event on the `ops_trace` queue to `services.tracing.delivery.TraceDeliveryService`: +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. @@ -92,3 +92,37 @@ Stored configurations keep the provider's fields at the top level, with `secret( 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. + +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/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/tracing/adapters.py b/api/services/tracing/adapters.py index 802454e55cf207..76976fda5b9c64 100644 --- a/api/services/tracing/adapters.py +++ b/api/services/tracing/adapters.py @@ -27,7 +27,6 @@ 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" -TRACE_QUEUE = "ops_trace" PAYLOAD_PREFIX = "trace_payloads" INLINE_EVENT_BYTES = 256 * 1024 # Node executions may be saved by another worker after their run row; give them time to land. @@ -50,7 +49,7 @@ def emit(self, event: TraceEvent) -> None: else: envelope = {"event": encoded} countdown = WORKFLOW_SETTLE_SECONDS if event.payload.kind == "workflow_run" else 0 - celery_app.send_task(DELIVER_TASK, args=(envelope,), queue=TRACE_QUEUE, countdown=countdown) + celery_app.send_task(DELIVER_TASK, args=(envelope,), countdown=countdown) class StoredPolicyTraceScopes(ExecutionTraceScopes): @@ -89,7 +88,7 @@ class CeleryExportRetryQueue(ExportRetryQueue): 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,), queue=TRACE_QUEUE, countdown=delay_seconds) + celery_app.send_task(RETRY_TASK, args=(key,), countdown=delay_seconds) class WorkspaceTracingSecretCipher(TracingSecretCipher): @@ -107,7 +106,7 @@ def mask(self, value: str) -> str: def schedule_agent_trace_export(run_id: str) -> None: - celery_app.send_task(AGENT_TRACE_EXPORT_TASK, args=(run_id,), queue=TRACE_QUEUE) + celery_app.send_task(AGENT_TRACE_EXPORT_TASK, args=(run_id,)) def send_to_enterprise_telemetry(record: BaseTraceInfo) -> None: diff --git a/api/tasks/application_trace_task.py b/api/tasks/application_trace_task.py index 6e1f6ae1b96007..e1ea6cbe05e9bd 100644 --- a/api/tasks/application_trace_task.py +++ b/api/tasks/application_trace_task.py @@ -1,15 +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 extensions.ext_application_services import application_services -@shared_task(queue="ops_trace") +@shared_task def export_application_trace(run_id: str) -> None: application_services().tracing.agent_traces.deliver(run_id) -@shared_task(queue="ops_trace") +@shared_task def reconcile_application_traces() -> None: application_services().tracing.agent_traces.reconcile() diff --git a/api/tasks/ops_trace_task.py b/api/tasks/ops_trace_task.py index 8dbfd4e38da9e2..cd78df32fb41bf 100644 --- a/api/tasks/ops_trace_task.py +++ b/api/tasks/ops_trace_task.py @@ -2,7 +2,8 @@ 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; -``sweep_trace_payloads`` removes any payload this task never sees. +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 logging diff --git a/api/tasks/trace_delivery_task.py b/api/tasks/trace_delivery_task.py index b45ac0029642f4..38dc69bf1fd989 100644 --- a/api/tasks/trace_delivery_task.py +++ b/api/tasks/trace_delivery_task.py @@ -1,4 +1,4 @@ -"""Celery entry points of the trace pipeline, on the ``ops_trace`` queue. +"""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 @@ -23,7 +23,6 @@ REDELIVERY_LIMIT = 10 PAYLOAD_RETENTION = timedelta(days=2) -_LEGACY_PAYLOAD_PREFIX = "ops_trace/" _PAYLOAD_DAY = re.compile(rf"{PAYLOAD_PREFIX}/(\d{{4}}-\d{{2}}-\d{{2}})/") @@ -48,7 +47,7 @@ def _delete_quietly(key: str) -> None: logger.warning("Trace payload cleanup failed: key=%s", key, exc_info=True) -@shared_task(queue="ops_trace", bind=True, max_retries=REDELIVERY_LIMIT) +@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 @@ -84,7 +83,7 @@ def deliver_trace_event(self, envelope: dict[str, Any]) -> None: _delete_quietly(payload_key) -@shared_task(queue="ops_trace") +@shared_task def retry_trace_export(job_key: str) -> None: try: job = ExportJob.model_validate_json(storage.load(job_key)) @@ -98,18 +97,16 @@ def retry_trace_export(job_key: str) -> None: _delete_quietly(job_key) -@shared_task(queue="ops_trace") +@shared_task def sweep_trace_payloads() -> None: - """Delete payloads a crashed delivery left behind, and payloads of the previous pipeline.""" + """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. - for prefix in (f"{PAYLOAD_PREFIX}/", _LEGACY_PAYLOAD_PREFIX): - try: - keys = storage.scan(prefix, files=True) - except (FileNotFoundError, NotImplementedError): - continue - for key in keys: - if not key.startswith(prefix): - continue - if prefix == _LEGACY_PAYLOAD_PREFIX or _expired(key, cutoff): - _delete_quietly(key) + 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/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/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/configs/test_app_config.py b/api/tests/unit_tests/configs/test_app_config.py index 35bd1438ddb153..1238fb0294a379 100644 --- a/api/tests/unit_tests/configs/test_app_config.py +++ b/api/tests/unit_tests/configs/test_app_config.py @@ -20,6 +20,12 @@ def test_ops_trace_export_attempts_must_be_positive() -> None: OpsTraceConfig(OPS_TRACE_EXPORT_MAX_ATTEMPTS=0) +@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( ("value", "expected"), [ @@ -113,6 +119,7 @@ def test_dify_config() -> None: assert config.APP_MAX_EXECUTION_TIME == 3600 assert config.WORKFLOW_MAX_EXECUTION_TIME == 3600 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 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/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/tracing/test_adapters.py b/api/tests/unit_tests/services/tracing/test_adapters.py index 5afba1f01b2a70..b0404a5ec9fc86 100644 --- a/api/tests/unit_tests/services/tracing/test_adapters.py +++ b/api/tests/unit_tests/services/tracing/test_adapters.py @@ -12,7 +12,6 @@ DELIVER_TASK, INLINE_EVENT_BYTES, RETRY_TASK, - TRACE_QUEUE, WORKFLOW_SETTLE_SECONDS, CeleryExportRetryQueue, CeleryTraceEmitter, @@ -31,10 +30,10 @@ class _Broker: def __init__(self) -> None: - self.sent: list[tuple[str, tuple[object, ...], str]] = [] + self.sent: list[tuple[str, tuple[object, ...], str | None]] = [] self.countdowns: list[int] = [] - def send_task(self, name: str, *, args: tuple[object, ...], queue: str, countdown: int = 0) -> None: + def send_task(self, name: str, *, args: tuple[object, ...], queue: str | None = None, countdown: int = 0) -> None: self.sent.append((name, args, queue)) self.countdowns.append(countdown) @@ -102,7 +101,7 @@ def test_small_events_travel_inline(broker: _Broker, storage: _Storage) -> None: CeleryTraceEmitter().emit(event) ((name, (envelope,), queue),) = broker.sent - assert (name, queue) == (DELIVER_TASK, TRACE_QUEUE) + assert (name, queue) == (DELIVER_TASK, None) assert isinstance(envelope, dict) assert TraceEvent.model_validate_json(envelope["event"]) == event assert storage.files == {} @@ -169,7 +168,7 @@ def test_export_retries_persist_the_trace_and_run_after_their_delay(broker: _Bro CeleryExportRetryQueue().schedule(job, delay_seconds=20) ((name, (key,), queue),) = broker.sent - assert (name, queue, broker.countdowns) == (RETRY_TASK, TRACE_QUEUE, [20]) + assert (name, queue, broker.countdowns) == (RETRY_TASK, None, [20]) assert isinstance(key, str) assert key.startswith("trace_payloads/") assert "/exports/" in key @@ -179,7 +178,7 @@ def test_export_retries_persist_the_trace_and_run_after_their_delay(broker: _Bro def test_agent_trace_exports_run_on_the_trace_queue(broker: _Broker) -> None: schedule_agent_trace_export("run-1") - assert broker.sent == [(AGENT_TRACE_EXPORT_TASK, ("run-1",), TRACE_QUEUE)] + assert broker.sent == [(AGENT_TRACE_EXPORT_TASK, ("run-1",), None)] @pytest.mark.parametrize( diff --git a/api/tests/unit_tests/tasks/test_trace_delivery_task.py b/api/tests/unit_tests/tasks/test_trace_delivery_task.py index 9e70a12ba7d04d..16b8e47b8b708a 100644 --- a/api/tests/unit_tests/tasks/test_trace_delivery_task.py +++ b/api/tests/unit_tests/tasks/test_trace_delivery_task.py @@ -262,7 +262,7 @@ def test_unreadable_export_jobs_are_dropped(delivery: _Delivery, storage: _Stora assert storage.files == {} -def test_sweep_removes_expired_and_legacy_payloads_only(storage: _Storage) -> None: +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") @@ -271,4 +271,6 @@ def test_sweep_removes_expired_and_legacy_payloads_only(storage: _Storage) -> No sweep_trace_payloads.run() - assert sorted(storage.files) == sorted([fresh, "trace_payloads/not-a-date/x.json", "other/keep.json"]) + 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/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/README.md b/docker/README.md index e4c9af0d2b072c..fe7b7c92f7bfb8 100644 --- a/docker/README.md +++ b/docker/README.md @@ -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 diff --git a/docker/envs/core-services/shared.env.example b/docker/envs/core-services/shared.env.example index e97c2709ef7daf..dcd033c63065f9 100644 --- a/docker/envs/core-services/shared.env.example +++ b/docker/envs/core-services/shared.env.example @@ -114,6 +114,8 @@ ENABLE_REQUEST_LOGGING=False # Application-level Agent tracing; also enable the collector in dify-agent.env. OPS_AGENT_TRACE_ENABLED=false ENABLE_APPLICATION_TRACE_RECONCILE_TASK=true +# 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. @@ -137,7 +139,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= From 53a74c66e8df6a535a15018d3000f3eeaffb1039 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Sat, 31 Oct 2026 01:35:30 +0800 Subject: [PATCH 12/17] fix(tracing): validate bound execution interfaces Replace unchecked casts of Flask extension bindings with runtime protocol checks. Missing or invalid scope and Agent hook bindings remain unavailable so tracing cannot break execution, and the PR complies with the no-new-cast CI guard. --- api/extensions/tracing.py | 10 +++++++--- api/tests/unit_tests/extensions/test_tracing.py | 13 +++++++++++++ 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/api/extensions/tracing.py b/api/extensions/tracing.py index 93e60916e43063..d5273e28f8e6ee 100644 --- a/api/extensions/tracing.py +++ b/api/extensions/tracing.py @@ -7,7 +7,7 @@ import logging from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Literal, Protocol, cast +from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable from dify_trace_sdk.emitter import TraceScope from dify_trace_sdk.events import TraceCorrelation @@ -28,6 +28,7 @@ _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.""" @@ -46,6 +47,7 @@ def for_app( ) -> TraceScope | None: ... +@runtime_checkable class ApplicationTracingHooks(Protocol): """Agent trace collection lifecycle; it never replaces the trace pipeline.""" @@ -85,7 +87,8 @@ def bind_tracing(app: Flask, *, scopes: ExecutionTraceScopes, hooks: Application def _scopes() -> ExecutionTraceScopes | None: if not has_app_context(): return None - return cast(ExecutionTraceScopes | None, current_app.extensions.get(_SCOPES_KEY)) + scopes = current_app.extensions.get(_SCOPES_KEY) + return scopes if isinstance(scopes, ExecutionTraceScopes) else None def correlation_from_extras(extras: Mapping[str, Any]) -> TraceCorrelation: @@ -148,4 +151,5 @@ def application_tracing_hooks() -> ApplicationTracingHooks | None: """ if not has_app_context(): return None - return cast(ApplicationTracingHooks | None, current_app.extensions.get(_HOOKS_KEY)) + hooks = current_app.extensions.get(_HOOKS_KEY) + return hooks if isinstance(hooks, ApplicationTracingHooks) else None diff --git a/api/tests/unit_tests/extensions/test_tracing.py b/api/tests/unit_tests/extensions/test_tracing.py index 64ee89e7f73309..cec06ffa4e581c 100644 --- a/api/tests/unit_tests/extensions/test_tracing.py +++ b/api/tests/unit_tests/extensions/test_tracing.py @@ -171,3 +171,16 @@ def test_agent_hooks_are_the_ones_the_composition_root_bound() -> None: 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 From 966a5bdddd120a6b0d01c4c41b692fe091cc06ed Mon Sep 17 00:00:00 2001 From: -LAN- Date: Sat, 31 Oct 2026 02:41:51 +0800 Subject: [PATCH 13/17] refactor(docker): keep env example limited to deployment inputs Keep only explicit Compose service selection in the root env template. Load Docker-specific connection and startup defaults before optional service files and local overrides, and retain the optional configuration catalog. Preserve existing env files, comments, multiline values, and permissions during synchronization. Append only missing inputs so removing defaults from the template cannot delete operator settings. --- AGENTS.md | 2 +- .../contracts/configs/test_docker_env_sync.py | 56 ++++ docker/.env.example | 314 +----------------- docker/README.md | 27 +- docker/dify-env-sync.py | 81 ++--- docker/dify-env-sync.sh | 129 ++----- docker/docker-compose-template.yaml | 7 +- docker/docker-compose.yaml | 7 +- .../envs/core-services/dify-agent.env.example | 4 + docker/envs/core-services/shared.defaults | 28 ++ docker/envs/core-services/shared.env.example | 5 +- docker/envs/core-services/web.env.example | 2 + 12 files changed, 187 insertions(+), 475 deletions(-) create mode 100644 api/tests/e2e_tests/contracts/configs/test_docker_env_sync.py create mode 100644 docker/envs/core-services/shared.defaults 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/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/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 fe7b7c92f7bfb8..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 @@ -275,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`. @@ -300,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. @@ -318,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..18e6aea6abec7c 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 # ================================================================ @@ -107,8 +107,8 @@ detect_differences() { 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" @@ -261,123 +261,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 +337,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 } @@ -446,7 +373,7 @@ main() { # Detect differences detect_differences - # Detect removed variables (before sync) + # Report additional local overrides (before sync) detect_removed_variables # Synchronize environment file 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 dcd033c63065f9..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 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= From 5418605ad4a7c6697b0abb3724213018c0faf94b Mon Sep 17 00:00:00 2001 From: -LAN- Date: Sat, 31 Oct 2026 02:53:40 +0800 Subject: [PATCH 14/17] fix(tracing): preserve batch types and scope trace writes by tenant Keep AgentTraceBatch and AgentTraceSpan typed throughout collection while preserving canonical payload hashes and privacy filtering. Require tenant IDs for locked trace updates and propagate the admitted tenant through collection, delivery, retries, and recovery. Batch node payload metadata lookups across nested workflow executions, resolve files within each execution tenant, and read each shared blob once after closing the database session. Cover tenant isolation, bounded query counts, and the SQL and LogStore paths. --- api/domain/tracing/trace_run_agg.py | 26 ++-- ..._api_workflow_node_execution_repository.py | 14 +++ .../api_workflow_node_execution_repository.py | 19 +++ .../application_trace_repository.py | 8 +- ..._api_workflow_node_execution_repository.py | 43 ++++++- api/repositories/trace_record_repository.py | 20 +-- api/services/tracing/collection_ports.py | 4 +- api/services/tracing/recovery.py | 2 +- api/services/tracing/service.py | 54 +++++---- .../domain/tracing/test_terminal_recovery.py | 2 +- .../domain/tracing/test_trace_run_agg.py | 17 ++- .../test_application_trace_repository.py | 54 +++++---- ..._api_workflow_node_execution_repository.py | 114 +++++++++++++++++- .../services/tracing/test_recovery.py | 8 +- .../services/tracing/test_runtime_bridge.py | 12 +- .../services/tracing/test_service.py | 12 +- 16 files changed, 300 insertions(+), 109 deletions(-) diff --git a/api/domain/tracing/trace_run_agg.py b/api/domain/tracing/trace_run_agg.py index f836154cc726e4..edacb58cf979bb 100644 --- a/api/domain/tracing/trace_run_agg.py +++ b/api/domain/tracing/trace_run_agg.py @@ -15,6 +15,7 @@ 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 @@ -328,11 +329,10 @@ def open_invocation( self.invocations[identifier] = invocation return invocation - def ingest(self, payload: Mapping[str, Any], *, max_spans: int, max_bytes: int, now: datetime) -> int: + 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. - ``payload`` is the batch's validated JSON form. A batch ID replays only - with the same payload, even after the run closed. + 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. @@ -341,13 +341,13 @@ def ingest(self, payload: Mapping[str, Any], *, max_spans: int, max_bytes: int, """ # Hash the original validated payload. Sanitization cannot turn two # conflicting producer payloads into a falsely successful replay. - canonical = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + canonical = json.dumps(batch.model_dump(mode="json"), sort_keys=True, separators=(",", ":")).encode() digest = hashlib.sha256(canonical).hexdigest() - invocation = self.invocations.get(str(payload["invocation_id"])) + invocation = self.invocations.get(str(batch.invocation_id)) if invocation is None: raise TraceRunNotFoundError("tracing_invocation_not_found") - batch_id = str(payload["batch_id"]) - spans: Sequence[Mapping[str, Any]] = payload["spans"] + 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: @@ -355,11 +355,11 @@ def ingest(self, payload: Mapping[str, Any], *, max_spans: int, max_bytes: int, 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): + 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 = payload["final"] is True + final = batch.final self.new_batches.append( TraceBatch( id=batch_id, @@ -377,12 +377,12 @@ def ingest(self, payload: Mapping[str, Any], *, max_spans: int, max_bytes: int, self.invocations[invocation.id] = replace(invocation, finalized=True, finalized_at=now) return len(spans) - def _stored_span(self, span: Mapping[str, Any]) -> dict[str, Any]: - stored = dict(span) - stored["attributes"] = safe_attributes(span["attributes"], capture_content=self.capture_content) + 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 + stored["status_message"] = "Agent operation failed" if span.status == "error" else None return stored def capture_terminal(self, trace: Trace, *, now: datetime) -> None: 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 24e8b125464967..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 ( @@ -174,6 +175,19 @@ def load_full_process_data(self, execution: WorkflowNodeExecutionModel) -> Mappi """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/repositories/api_workflow_node_execution_repository.py b/api/repositories/api_workflow_node_execution_repository.py index 20be54e9d522ad..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: """ @@ -192,6 +201,16 @@ def load_full_outputs(self, execution: WorkflowNodeExecutionModel) -> Mapping[st """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/application_trace_repository.py b/api/repositories/application_trace_repository.py index 171fcf8affa225..26dae5de587f38 100644 --- a/api/repositories/application_trace_repository.py +++ b/api/repositories/application_trace_repository.py @@ -448,9 +448,13 @@ def load(self, run_id: str) -> TraceRunAgg | None: @override @contextmanager - def change(self, run_id: str) -> Generator[LockedTraceRun | None]: + 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).with_for_update()) + 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 diff --git a/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py b/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py index a95afaf258d730..632cc2e55e3cea 100644 --- a/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py @@ -8,9 +8,10 @@ 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 @@ -25,6 +26,7 @@ ) from repositories.api_workflow_node_execution_repository import ( DifyAPIWorkflowNodeExecutionRepository, + WorkflowNodeExecutionPayloads, WorkflowNodeExecutionSnapshot, WorkflowNodeExecutionSnapshotRow, ) @@ -81,14 +83,43 @@ def _load_full_payload( 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: - key = session.scalar( - select(UploadFile.key).where( - UploadFile.id == offload.file_id, UploadFile.tenant_id == execution.tenant_id + 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), ) - assert key is not None, f"UploadFile with id {offload.file_id} should exist in the execution's tenant" - return json.loads(storage.load(key)) + return result @override def get_node_last_execution( diff --git a/api/repositories/trace_record_repository.py b/api/repositories/trace_record_repository.py index cc4a7c8f07f4f8..821505601888b5 100644 --- a/api/repositories/trace_record_repository.py +++ b/api/repositories/trace_record_repository.py @@ -41,7 +41,10 @@ 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 +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 @@ -381,7 +384,7 @@ def _workflow(self, event: TraceEvent, payload: WorkflowRunCompleted) -> Workflo def _node_executions(self, run: WorkflowRun) -> tuple[Sequence[WorkflowNodeExecution], Mapping[str, str]]: repository = self._node_executions_repository - executions: list[WorkflowNodeExecution] = [] + rows: list[WorkflowNodeExecutionModel] = [] tool_parents: dict[str, str] = {} seen: set[str] = set() with tracing_store_errors(): @@ -395,7 +398,7 @@ def _node_executions(self, run: WorkflowRun) -> tuple[Sequence[WorkflowNodeExecu if row.id in seen or row.status == WorkflowNodeExecutionStatus.PAUSED: continue seen.add(row.id) - executions.append(self._node_execution(row, repository)) + rows.append(row) if parent_id is not None: tool_parents[row.id] = parent_id if row.node_type == BuiltinNodeTypes.TOOL: @@ -406,11 +409,12 @@ def _node_executions(self, run: WorkflowRun) -> tuple[Sequence[WorkflowNodeExecu run.tenant_id, run.id, row.node_execution_id or row.id ) ) - return executions, tool_parents + 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, repository: DifyAPIWorkflowNodeExecutionRepository + row: WorkflowNodeExecutionModel, payloads: WorkflowNodeExecutionPayloads ) -> WorkflowNodeExecution: return WorkflowNodeExecution( id=row.id, @@ -422,9 +426,9 @@ def _node_execution( node_id=row.node_id, node_type=row.node_type, title=row.title, - inputs=repository.load_full_inputs(row), - process_data=repository.load_full_process_data(row), - outputs=repository.load_full_outputs(row), + inputs=payloads.inputs, + process_data=payloads.process_data, + outputs=payloads.outputs, status=WorkflowNodeExecutionStatus(row.status), error=row.error, elapsed_time=row.elapsed_time, diff --git a/api/services/tracing/collection_ports.py b/api/services/tracing/collection_ports.py index d25fedd92d6377..473edcd678cc9c 100644 --- a/api/services/tracing/collection_ports.py +++ b/api/services/tracing/collection_ports.py @@ -41,10 +41,10 @@ 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) -> AbstractContextManager[LockedTraceRun | None]: + 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. The lock is held until the block + ``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. """ diff --git a/api/services/tracing/recovery.py b/api/services/tracing/recovery.py index 18e69ff0c16234..562207b463a3a1 100644 --- a/api/services/tracing/recovery.py +++ b/api/services/tracing/recovery.py @@ -120,6 +120,6 @@ def recover_terminal( logger.warning("Application trace recovery failed: run_id=%s error_type=%s", run.id, type(error).__name__) return False finally: - with runs.change(run.id) as locked: + 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/service.py b/api/services/tracing/service.py index 39c1a8fcfd827a..ae3bd00ef3661e 100644 --- a/api/services/tracing/service.py +++ b/api/services/tracing/service.py @@ -265,7 +265,7 @@ def prepare_invocation( ) if run is None or not run.capture_agent: return None - with _collection_errors(), self._runs.change(run.id) as locked: + 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( @@ -289,13 +289,16 @@ def ingest(self, batch: AgentTraceBatch) -> int: TraceCollectionLimitError: The batch exceeds the run's limits. """ run_id = str(batch.run_id) - with _collection_errors(), self._runs.change(run_id) as locked: + 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.model_dump(mode="json"), + batch, max_spans=self._settings.max_spans, max_bytes=self._settings.max_bytes, now=naive_utc_now(), @@ -377,9 +380,9 @@ def _schedule(self, run_id: str) -> None: except Exception: logger.warning("Application trace is persisted; scheduling deferred: run_id=%s", run_id) - def _change(self, run_id: str, change: Callable[[TraceRunAgg], object]) -> None: + 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) as locked: + with self._runs.change(run.id, tenant_id=run.tenant_id) as locked: if locked is not None: change(locked.run) @@ -512,7 +515,7 @@ def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: trace = self._builder.build(info) if trace is None: raise TraceCollectionError("tracing_terminal_not_supported") - self._change(run.id, lambda current: current.capture_terminal(trace, now=naive_utc_now())) + self._change(run, lambda current: current.capture_terminal(trace, now=naive_utc_now())) self._schedule(run.id) return True @@ -521,11 +524,11 @@ def _settle_preview(self, run: TraceRunAgg) -> None: return 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.id, lambda current: current.settle_preview(node, now=naive_utc_now())) + self._change(run, lambda current: current.settle_preview(node, now=naive_utc_now())) - def _claim(self, run_id: str) -> tuple[TraceRunAgg, str] | None: + def _claim(self, run: TraceRunAgg) -> tuple[TraceRunAgg, str] | None: token = str(uuidv7()) - with self._runs.change(run_id) as locked: + 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, @@ -536,25 +539,26 @@ def _claim(self, run_id: str) -> tuple[TraceRunAgg, str] | None: return None return locked.run, token - def _renew_export_lease(self, run_id: str, token: str) -> bool: - with self._runs.change(run_id) as locked: + 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 = self._settings.export_lease_seconds / 3 def heartbeat() -> None: while not stopped.wait(interval): try: - if not self._renew_export_lease(run_id, token): + if not self._renew_export_lease(run, token): return except Exception as error: logger.warning( @@ -598,11 +602,11 @@ def deliver(self, run_id: str) -> None: if run is None: return self._settle_preview(run) - claim = self._claim(run_id) + 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 = run.exported_trace(self._run_query.batches(run.id)) @@ -613,20 +617,20 @@ def deliver(self, run_id: str) -> None: 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_id, lambda current: current.export_uncertain(token, now=naive_utc_now())) + self._change(run, lambda current: current.export_uncertain(token, now=naive_utc_now())) except ExportRetryableError: self._change( - run_id, + 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_id, token, error.code) + self._export_failed(run, token, error.code) except ProviderConfigUnreadableError: - self._export_failed(run_id, token, "tracing_config_unreadable") + self._export_failed(run, token, "tracing_config_unreadable") except ExportRejectedError: - self._export_failed(run_id, token, "provider_export_rejected") + 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", @@ -634,10 +638,10 @@ def deliver(self, run_id: str) -> None: run.provider, type(error).__name__, ) - self._export_failed(run_id, token, "provider_export_rejected") + self._export_failed(run, token, "provider_export_rejected") else: self._change( - run_id, + run, lambda current: current.export_accepted( token, provider_trace_id=result.trace_id, @@ -647,8 +651,8 @@ def deliver(self, run_id: str) -> None: ), ) - def _export_failed(self, run_id: str, token: str, error_code: str) -> None: - self._change(run_id, lambda current: current.export_failed(token, error_code, 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: @@ -794,7 +798,7 @@ def retry( if run is None: raise TraceCollectionNotFoundError("tracing_run_not_found") self._status_for_run(run, source_kind, source_id, node_execution_id) - with _collection_errors(), self._runs.change(run.id) as locked: + 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( diff --git a/api/tests/unit_tests/domain/tracing/test_terminal_recovery.py b/api/tests/unit_tests/domain/tracing/test_terminal_recovery.py index 376932d9e29b67..4f0674496c8c6e 100644 --- a/api/tests/unit_tests/domain/tracing/test_terminal_recovery.py +++ b/api/tests/unit_tests/domain/tracing/test_terminal_recovery.py @@ -55,7 +55,7 @@ def with_invocation(run: TraceRunAgg, *, finalized_at: datetime | None) -> Trace ) ], ) - run.ingest(payload.model_dump(mode="json"), max_spans=100, max_bytes=1_000_000, now=finalized_at) + run.ingest(payload, max_spans=100, max_bytes=1_000_000, now=finalized_at) return run 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 index 8fe1b1d90b659c..853502fbb2afe1 100644 --- a/api/tests/unit_tests/domain/tracing/test_trace_run_agg.py +++ b/api/tests/unit_tests/domain/tracing/test_trace_run_agg.py @@ -1,4 +1,3 @@ -from collections.abc import Mapping from dataclasses import replace from datetime import UTC, datetime, timedelta from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 @@ -58,8 +57,8 @@ def batch( batch_id: UUID | None = None, name: str = "Agent execution", status: str = "ok", -) -> dict[str, object]: - """The JSON form of a runtime batch, as the service hands it to the run.""" +) -> 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), @@ -79,10 +78,10 @@ def batch( events=[AgentTraceEvent(name="token", timestamp_ns=1_500_000_000)], ) ], - ).model_dump(mode="json") + ) -def ingest(run: TraceRunAgg, payload: Mapping[str, object], *, now: datetime = NOW) -> int: +def ingest(run: TraceRunAgg, payload: AgentTraceBatch, *, now: datetime = NOW) -> int: return run.ingest(payload, max_spans=100, max_bytes=1_000_000, now=now) @@ -277,7 +276,7 @@ def test_final_batch_is_stored_and_finalizes_its_invocation() -> None: assert ingest(run, payload) == 1 (stored,) = run.new_batches - assert (stored.id, stored.invocation_id, stored.final) == (payload["batch_id"], opened.id, True) + 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 @@ -304,7 +303,7 @@ def test_batch_id_replays_only_with_the_same_payload_and_invocation() -> None: other = invocation(run, "second") payload = batch(run, opened, final=False) ingest(run, payload) - batch_id = UUID(str(payload["batch_id"])) + 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")) @@ -361,7 +360,7 @@ def test_spans_must_belong_to_the_run_trace() -> None: ) with pytest.raises(TraceRunConflictError, match="tracing_trace_id_mismatch"): - ingest(run, foreign.model_dump(mode="json")) + ingest(run, foreign) assert run.new_batches == [] @staticmethod @@ -655,7 +654,7 @@ def report(run: TraceRunAgg, opened: TraceInvocation, spans: list[AgentTraceSpan payload = AgentTraceBatch( run_id=UUID(run.id), invocation_id=UUID(opened.id), batch_id=uuid4(), final=True, spans=spans ) - ingest(run, payload.model_dump(mode="json")) + ingest(run, payload) class TestExportedTrace: 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 068cd3d95fd6fa..86ebed894751c4 100644 --- a/api/tests/unit_tests/repositories/test_application_trace_repository.py +++ b/api/tests/unit_tests/repositories/test_application_trace_repository.py @@ -175,7 +175,7 @@ def seed(sessions: sessionmaker[Session], name: str) -> Collection: def open_invocation(collection: Collection, *, node_execution_id: str | None, call_key: str) -> TraceInvocation: - with collection.repository.change(collection.run.id) as locked: + 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 @@ -209,9 +209,9 @@ def ingest(collection: Collection, batch: AgentTraceBatch) -> int: 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) as locked: + with collection.repository.change(collection.run.id, tenant_id=collection.run.tenant_id) as locked: assert locked is not None - locked.run.ingest(batch.model_dump(mode="json"), max_spans=100, max_bytes=1_000_000, now=naive_utc_now()) + 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: @@ -236,7 +236,7 @@ def agent_turn(run: TraceRun, *, root_span_id: str) -> Trace: def capture_terminal(collection: Collection, trace: Trace) -> None: - with collection.repository.change(collection.run.id) as locked: + 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()) @@ -248,7 +248,7 @@ def terminal(collection: Collection) -> None: 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) as locked: + 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() @@ -612,7 +612,7 @@ def test_workflow_node_deep_link_is_scoped_to_durable_invocation(stored: StoredS terminal(case) token = claim(case) assert token is not None - with case.repository.change(run.id) as locked: + with case.repository.change(run.id, tenant_id=run.tenant_id) as locked: assert locked is not None locked.run.export_accepted( token, @@ -827,9 +827,9 @@ class PurgedBeforeLock(ApplicationTraceRepository): @override @contextmanager - def change(self, run_id: str) -> Generator[LockedTraceRun | None]: + def change(self, run_id: str, *, tenant_id: str) -> Generator[LockedTraceRun | None]: self.purge(retention_days=0, limit=100) - with super().change(run_id) as locked: + with super().change(run_id, tenant_id=tenant_id) as locked: yield locked @@ -1130,7 +1130,7 @@ def stored_fields(run: TraceRunAgg) -> dict[str, object]: def reloaded(collection: Collection) -> dict[str, object]: - with collection.repository.change(collection.run.id) as locked: + 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) @@ -1360,23 +1360,23 @@ def test_change_saves_root_fields_invocations_and_new_batches( case = seed(sqlite_session_factory, "saved") now = naive_utc_now() batches: list[TraceBatch] = [] - with case.repository.change(case.run.id) as locked: + 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.model_dump(mode="json"), max_spans=100, max_bytes=1_000_000, now=now) + 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) as locked: + 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.model_dump(mode="json"), max_spans=100, max_bytes=1_000_000, now=now) + 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 @@ -1385,14 +1385,14 @@ def test_change_saves_root_fields_invocations_and_new_batches( saved = stored_fields(locked.run) assert reloaded(case) == saved - with case.repository.change(case.run.id) as locked: + 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) as locked: + 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( @@ -1417,15 +1417,13 @@ def test_change_rolls_back_the_whole_block_when_it_raises(sqlite_session_factory before = reloaded(case) def change_then_break_a_rule() -> None: - with case.repository.change(case.run.id) as locked: + 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.model_dump(mode="json"), max_spans=100, max_bytes=1_000_000, now=naive_utc_now() - ) + 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() ) @@ -1464,7 +1462,7 @@ def test_change_tells_whether_the_owning_app_still_exists( case = seed(sqlite_session_factory, "owner") def owner_exists() -> bool: - with case.repository.change(case.run.id) as locked: + with case.repository.change(case.run.id, tenant_id=case.run.tenant_id) as locked: assert locked is not None return locked.owner_exists @@ -1478,10 +1476,20 @@ def owner_exists() -> bool: 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")) as locked: + with case.repository.change(identifier("missing-run"), tenant_id=case.run.tenant_id) as locked: assert locked is None @staticmethod @@ -1546,14 +1554,14 @@ def test_runs_read_invocations_by_attempt_then_creation_and_batches_by_arrival( case = seed(sqlite_session_factory, "collected") now = naive_utc_now() agent_id = identifier("agent") - with case.repository.change(case.run.id) as locked: + 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.model_dump(mode="json"), max_spans=100, max_bytes=1_000_000, now=now) + 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) 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 7d0829a38ee63e..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,11 +1,12 @@ 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, @@ -31,11 +32,13 @@ @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", @@ -54,7 +57,11 @@ def test_full_payloads_load_owned_files_after_closing_the_database_session( with sqlite_session_factory.begin() as session: session.add(file) execution = WorkflowNodeExecutionModel( - tenant_id="tenant", inputs='{"truncated":true}', outputs='{"truncated":true}', process_data='{"truncated":true}' + id="execution", + tenant_id="tenant", + inputs='{"truncated":true}', + outputs='{"truncated":true}', + process_data='{"truncated":true}', ) execution.offload_data = [] repository = DifyAPISQLAlchemyWorkflowNodeExecutionRepository(sqlite_session_factory) @@ -63,7 +70,18 @@ def test_full_payloads_load_owned_files_after_closing_the_database_session( ExecutionOffLoadType.OUTPUTS: repository.load_full_outputs, ExecutionOffLoadType.PROCESS_DATA: repository.load_full_process_data, }[payload_type] - assert load(execution) == {"truncated": True} + + 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 @@ -85,11 +103,97 @@ def load_storage(key: str) -> bytes: try: if foreign_tenant: with pytest.raises(AssertionError, match="execution's tenant"): - load(execution) + read_payload() else: - assert load(execution) == {"complete": "payload"} + 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]) diff --git a/api/tests/unit_tests/services/tracing/test_recovery.py b/api/tests/unit_tests/services/tracing/test_recovery.py index 89fbcbd3e9a714..7b8b19661e94fc 100644 --- a/api/tests/unit_tests/services/tracing/test_recovery.py +++ b/api/tests/unit_tests/services/tracing/test_recovery.py @@ -106,7 +106,7 @@ class Collection: run: TraceRun def invocation(self, name: str) -> TraceInvocation: - with self.repository.change(self.run.id) as locked: + 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 @@ -309,15 +309,15 @@ def batch_for(run: TraceRun, invocation: TraceInvocation, *, final: bool, span_i 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)) as locked: + 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.model_dump(mode="json"), max_spans=100, max_bytes=1_000_000, now=naive_utc_now()) + 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) as locked: + 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() 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 b3c8c810c71369..9cf24c29429712 100644 --- a/api/tests/unit_tests/services/tracing/test_runtime_bridge.py +++ b/api/tests/unit_tests/services/tracing/test_runtime_bridge.py @@ -326,16 +326,20 @@ def trace_event(run: TraceRun, payload: TracePayload) -> TraceEvent: def capture(traces: ApplicationTraceRepository, run_id: str, trace: Trace) -> None: """Keep a business terminal in its run without scheduling the run's export.""" - with traces.change(run_id) as locked: + 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.""" - with traces.change(str(batch.run_id)) as locked: + 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.model_dump(mode="json"), max_spans=max_spans, max_bytes=max_bytes, now=naive_utc_now()) + locked.run.ingest(batch, max_spans=max_spans, max_bytes=max_bytes, now=naive_utc_now()) def terminal(collection: Collection) -> None: @@ -776,7 +780,7 @@ def test_mixed_workflow_exports_normal_llm_and_agent_calls_without_double_usage( ) node_record.created_at = row.created_at session.add(node_record) - with case.collection.repository.change(managed.id) as locked: + 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" diff --git a/api/tests/unit_tests/services/tracing/test_service.py b/api/tests/unit_tests/services/tracing/test_service.py index d6e26367c291c8..aa30f07f9b7678 100644 --- a/api/tests/unit_tests/services/tracing/test_service.py +++ b/api/tests/unit_tests/services/tracing/test_service.py @@ -86,7 +86,7 @@ class _ApplicationTraceCollectionCollection: @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) as locked: + with self.repository.change(self.run.id, tenant_id=self.run.tenant_id) as locked: assert locked is not None yield locked.run @@ -445,7 +445,7 @@ def _claim_at( ) -> 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) as locked: + 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 @@ -1926,9 +1926,9 @@ def __init__(self, sessions: sessionmaker[Session], *, transient_failure: bool) @override @contextmanager - def change(self, run_id: str) -> Generator[LockedTraceRun | None]: + def change(self, run_id: str, *, tenant_id: str) -> Generator[LockedTraceRun | None]: renewal = False - with super().change(run_id) as locked: + 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 @@ -1945,8 +1945,8 @@ class _UnavailableCompletionRepository(ApplicationTraceRepository): @override @contextmanager - def change(self, run_id: str) -> Generator[LockedTraceRun | None]: - with super().change(run_id) as locked: + 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: From 2cac589a81f17c0f38fa35bf03039ec57ecf35d1 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Sat, 31 Oct 2026 02:53:44 +0800 Subject: [PATCH 15/17] fix(docker): satisfy environment sync script style checks Separate command substitutions from local declarations so command failures remain visible, use an explicit no-op for file truncation, and restore the required final newline. --- docker/dify-env-sync.sh | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/docker/dify-env-sync.sh b/docker/dify-env-sync.sh index 18e6aea6abec7c..7d64c00011a989 100755 --- a/docker/dify-env-sync.sh +++ b/docker/dify-env-sync.sh @@ -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,7 +104,8 @@ 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 @@ -111,10 +113,11 @@ detect_differences() { 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:]]*$/ && /=/) { @@ -351,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" @@ -389,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 From 44f72fa9988e2214153c9f317231075ac644ec1a Mon Sep 17 00:00:00 2001 From: -LAN- Date: Sat, 31 Oct 2026 03:12:32 +0800 Subject: [PATCH 16/17] fix(tracing): isolate producer I/O and application imports Publish trace events through a bounded background executor while preserving payload snapshots and the caller Celery app. Drain accepted events during shutdown, reset publishers after forks, and log publication failures or overflow without blocking producers on storage or broker I/O. Move shared workflow retry history and Agent node identification into pure domain modules and update all consumers. Forbid legacy core imports in the tracing application boundary and cover cold imports, blocked publication, capacity, fork isolation, and shutdown with focused regression checks. --- api/.importlinter | 1 + api/core/workflow/node_factory.py | 2 +- .../workflow/nodes/agent_v2/validators.py | 2 +- api/domain/workflow/__init__.py | 1 + .../workflow/agent_node.py} | 2 + .../app => domain}/workflow/retry_history.py | 2 +- ...ngs_when_app_published_workflow_updated.py | 2 +- api/providers/trace/README.md | 4 + api/services/agent/composer_candidates.py | 2 +- api/services/agent/dsl_service.py | 2 +- api/services/tracing/adapters.py | 86 +++++- api/services/tracing/canonical.py | 2 +- api/services/tracing/collection.py | 2 +- .../workflow/node_output_inspector_service.py | 2 +- .../workflow_node_execution_trace_service.py | 2 +- api/services/workflow_persistence.py | 2 +- api/services/workflow_service.py | 2 +- .../workflow/nodes/agent_v2/test_entities.py | 9 + .../workflow/test_agent_node.py} | 9 +- ...t_workflow_node_execution_trace_service.py | 2 +- .../services/test_workflow_persistence.py | 2 +- .../services/tracing/test_adapters.py | 255 ++++++++++++++---- .../services/tracing/test_canonical.py | 34 +++ 23 files changed, 346 insertions(+), 83 deletions(-) create mode 100644 api/domain/workflow/__init__.py rename api/{core/workflow/nodes/agent_v2/discriminator.py => domain/workflow/agent_node.py} (92%) rename api/{core/app => domain}/workflow/retry_history.py (88%) create mode 100644 api/tests/unit_tests/core/workflow/nodes/agent_v2/test_entities.py rename api/tests/unit_tests/{core/workflow/nodes/agent_v2/test_discriminator.py => domain/workflow/test_agent_node.py} (69%) diff --git a/api/.importlinter b/api/.importlinter index 11069a69fbf6b9..6b506a50edba1e 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -1341,6 +1341,7 @@ forbidden_modules = celery configs controllers + core enterprise extensions flask 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/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/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/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/providers/trace/README.md b/api/providers/trace/README.md index 984c113b3b09da..4e445daceb0f1b 100644 --- a/api/providers/trace/README.md +++ b/api/providers/trace/README.md @@ -84,6 +84,10 @@ Stored configurations keep the provider's fields at the top level, with `secret( - 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. 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/tracing/adapters.py b/api/services/tracing/adapters.py index 76976fda5b9c64..b4d4a024b5f906 100644 --- a/api/services/tracing/adapters.py +++ b/api/services/tracing/adapters.py @@ -1,16 +1,22 @@ """Runtime adapters of the trace pipeline over Celery, storage, workspace keys, and enterprise telemetry. -Events go straight to the delivery worker; large ones travel through storage -instead of the broker. Every payload key carries its date so the sweeper can -expire payloads that a crashed delivery left behind. +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. """ -from collections.abc import Sequence +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 @@ -29,9 +35,12 @@ 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.""" @@ -39,17 +48,68 @@ def payload_key(day: datetime, kind: str, identifier: str) -> str: 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: - encoded = event.model_dump_json() - if len(encoded.encode()) > INLINE_EVENT_BYTES: - key = payload_key(event.occurred_at, "events", event.event_id) - storage.save(key, encoded.encode()) - envelope = {"payload_key": key} - else: - envelope = {"event": encoded} - countdown = WORKFLOW_SETTLE_SECONDS if event.payload.kind == "workflow_run" else 0 - celery_app.send_task(DELIVER_TASK, args=(envelope,), countdown=countdown) + 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): diff --git a/api/services/tracing/canonical.py b/api/services/tracing/canonical.py index 0aa1d484b9d68a..0f5bcb2ffd4253 100644 --- a/api/services/tracing/canonical.py +++ b/api/services/tracing/canonical.py @@ -19,7 +19,7 @@ from dify_trace_sdk.model import Cost, Span, SpanKind, SpanRef, SpanStatus, TokenUsage, Trace from pydantic import JsonValue -from core.app.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY +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 diff --git a/api/services/tracing/collection.py b/api/services/tracing/collection.py index 5b31a216d59b51..f0280fb3fe8280 100644 --- a/api/services/tracing/collection.py +++ b/api/services/tracing/collection.py @@ -13,9 +13,9 @@ from enum import StrEnum from typing import Any -from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data 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 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 1cf938839b8a4f..6707c010ebd482 100644 --- a/api/services/workflow_persistence.py +++ b/api/services/workflow_persistence.py @@ -13,7 +13,6 @@ 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.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from core.tools.workflow_as_tool.repository import WorkflowToolSource from core.workflow.node_execution import WorkflowNodeExecution as DifyWorkflowNodeExecution @@ -21,6 +20,7 @@ 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 ( diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py index 1ee77552ec5fb0..9777702d0ece3f 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/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_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/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 ea38df25573781..1d3bf637c4a6dd 100644 --- a/api/tests/unit_tests/services/test_workflow_persistence.py +++ b/api/tests/unit_tests/services/test_workflow_persistence.py @@ -15,7 +15,6 @@ 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.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,6 +23,7 @@ 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 diff --git a/api/tests/unit_tests/services/tracing/test_adapters.py b/api/tests/unit_tests/services/tracing/test_adapters.py index b0404a5ec9fc86..1afaf8a09318e3 100644 --- a/api/tests/unit_tests/services/tracing/test_adapters.py +++ b/api/tests/unit_tests/services/tracing/test_adapters.py @@ -1,16 +1,29 @@ +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, @@ -28,24 +41,6 @@ ) -class _Broker: - def __init__(self) -> None: - self.sent: list[tuple[str, tuple[object, ...], str | None]] = [] - self.countdowns: list[int] = [] - - def send_task(self, name: str, *, args: tuple[object, ...], queue: str | None = None, countdown: int = 0) -> None: - self.sent.append((name, args, queue)) - self.countdowns.append(countdown) - - -class _Storage: - def __init__(self) -> None: - self.files: dict[str, bytes] = {} - - def save(self, key: str, data: bytes) -> None: - self.files[key] = data - - class _Recorder: def __init__(self) -> None: self.events: list[TraceEvent] = [] @@ -61,17 +56,31 @@ def policy(self, *, tenant_id: str, app_id: str) -> AppTraceSettings | None: @pytest.fixture -def broker(monkeypatch: pytest.MonkeyPatch) -> _Broker: - broker = _Broker() - monkeypatch.setattr(adapters, "celery_app", broker) - return broker +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) -> _Storage: - storage = _Storage() - monkeypatch.setattr(adapters, "storage", storage) - return storage +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: @@ -95,51 +104,66 @@ def _event(query: str) -> TraceEvent: ) -def test_small_events_travel_inline(broker: _Broker, storage: _Storage) -> None: +def test_small_events_travel_inline(broker: Celery, storage: OpenDALStorage) -> None: event = _event("hello") - CeleryTraceEmitter().emit(event) + emitter = CeleryTraceEmitter() + emitter.emit(event) + emitter.shutdown() - ((name, (envelope,), queue),) = broker.sent - assert (name, queue) == (DELIVER_TASK, None) + (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 storage.files == {} + assert not storage.exists("trace_payloads") -def test_large_events_travel_through_storage(broker: _Broker, storage: _Storage) -> None: +def test_large_events_travel_through_storage(broker: Celery, storage: OpenDALStorage) -> None: event = _event("x" * INLINE_EVENT_BYTES) - CeleryTraceEmitter().emit(event) + emitter = CeleryTraceEmitter() + emitter.emit(event) + emitter.shutdown() - ((_, (envelope,), _),) = broker.sent + (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.files["trace_payloads/2026-03-04/events/event-1.json"]) == event + 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: _Broker) -> None: - CeleryTraceEmitter().emit(_event("中" * (INLINE_EVENT_BYTES // 2))) +def test_inline_size_is_measured_in_bytes(broker: Celery) -> None: + emitter = CeleryTraceEmitter() + emitter.emit(_event("中" * (INLINE_EVENT_BYTES // 2))) + emitter.shutdown() - ((_, (envelope,), _),) = broker.sent + (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: _Broker) -> None: +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)} ) - CeleryTraceEmitter().emit(_event("hello")) - CeleryTraceEmitter().emit(workflow) + emitter = CeleryTraceEmitter() + emitter.emit(_event("hello")) + before = datetime.now(UTC) + emitter.emit(workflow) + emitter.shutdown() - assert broker.countdowns == [0, WORKFLOW_SETTLE_SECONDS] + 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: _Broker, storage: _Storage) -> None: +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", @@ -167,18 +191,153 @@ def test_export_retries_persist_the_trace_and_run_after_their_delay(broker: _Bro CeleryExportRetryQueue().schedule(job, delay_seconds=20) - ((name, (key,), queue),) = broker.sent - assert (name, queue, broker.countdowns) == (RETRY_TASK, None, [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.files[key]) == job + assert ExportJob.model_validate_json(storage.load_once(key)) == job -def test_agent_trace_exports_run_on_the_trace_queue(broker: _Broker) -> None: +def test_agent_trace_exports_run_on_the_trace_queue(broker: Celery) -> None: schedule_agent_trace_export("run-1") - assert broker.sent == [(AGENT_TRACE_EXPORT_TASK, ("run-1",), None)] + (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( diff --git a/api/tests/unit_tests/services/tracing/test_canonical.py b/api/tests/unit_tests/services/tracing/test_canonical.py index c4d9633dc73ad3..bf6b7cfcdcb7ee 100644 --- a/api/tests/unit_tests/services/tracing/test_canonical.py +++ b/api/tests/unit_tests/services/tracing/test_canonical.py @@ -1,5 +1,9 @@ +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 @@ -359,3 +363,33 @@ def test_records_without_a_provider_mapping_are_not_exported() -> 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 From e5ccad5d1a71edfa519b5bfc7e1766591fb5d858 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Sat, 31 Oct 2026 03:22:54 +0800 Subject: [PATCH 17/17] test(api): reset Redis lock state between unit tests Clear the shared Redis lock mock return value and side effects before each unit test. This prevents a failed lock acquisition in database migration tests from leaking into credit-accounting tests and causing order-dependent CI failures. --- api/tests/unit_tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/tests/unit_tests/conftest.py b/api/tests/unit_tests/conftest.py index ba969cda9d258b..b5cf36363aefbb 100644 --- a/api/tests/unit_tests/conftest.py +++ b/api/tests/unit_tests/conftest.py @@ -133,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()