opensource
/
dify
mirror of https://github.com/langgenius/dify.git
2
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactor(telemetry): migrate to type-safe enum-based event routing with centralized enterprise filtering 

Changes:
- Change TelemetryEvent.name from str to TraceTaskName enum for type safety
- Remove hardcoded trace_task_name_map from facade (no mapping needed)
- Add centralized enterprise-only filter in TelemetryFacade.emit()
- Rename is_telemetry_enabled() to is_enterprise_telemetry_enabled()
- Update all 11 call sites to pass TraceTaskName enum values
- Remove redundant enterprise guard from draft_trace.py
- Add unit tests for TelemetryFacade.emit() routing (6 tests)
- Add unit tests for TraceQueueManager telemetry guard (5 tests)
- Fix test fixture scoping issue for full test suite compatibility
- Fix tenant_id handling in agent tool callback handler

Benefits:
- 100% type-safe: basedpyright catches errors at compile time
- No string literals: eliminates entire class of typo bugs
- Single point of control: centralized filtering in facade
- All guards removed except facade
- Zero regressions: 4887 tests passing

Verification:
- make lint: PASS
- make type-check: PASS (0 errors, 0 warnings)
- pytest: 4887 passed, 8 skipped
This commit is contained in:
GareArc 2026-02-05 15:12:02 -08:00
parent ed222945aa
commit adadf1ec5f
No known key found for this signature in database
16 changed files with 502 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
api/core/agent/base_agent_runner.py
View File

@ -79,7 +79,7 @@ class BaseAgentRunner(AppRunner):
self.model_instance = model_instance
# init callback
self.agent_callback = DifyAgentCallbackHandler()
self.agent_callback = DifyAgentCallbackHandler(tenant_id=tenant_id)
# init dataset tools
hit_callback = DatasetIndexToolCallbackHandler(
queue_manager=queue_manager,

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

@ -63,7 +63,7 @@ from core.base.tts import AppGeneratorTTSPublisher, AudioTrunk
from core.model_runtime.entities.llm_entities import LLMUsage
from core.model_runtime.utils.encoders import jsonable_encoder
from core.ops.ops_trace_manager import TraceQueueManager
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade, TraceTaskName
from core.workflow.enums import WorkflowExecutionStatus
from core.workflow.nodes import NodeType
from core.workflow.repositories.draft_variable_repository import DraftVariableSaverFactory
@ -834,7 +834,7 @@ class AdvancedChatAppGenerateTaskPipeline(GraphRuntimeStateSupport):
if trace_manager:
TelemetryFacade.emit(
TelemetryEvent(
name="message",
name=TraceTaskName.MESSAGE_TRACE,
context=TelemetryContext(
tenant_id=self._application_generate_entity.app_config.tenant_id,
app_id=self._application_generate_entity.app_config.app_id,

4
api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py
View File

@ -55,7 +55,7 @@ from core.model_runtime.model_providers.__base.large_language_model import Large
from core.ops.ops_trace_manager import TraceQueueManager
from core.prompt.utils.prompt_message_util import PromptMessageUtil
from core.prompt.utils.prompt_template_parser import PromptTemplateParser
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade, TraceTaskName
from events.message_event import message_was_created
from extensions.ext_database import db
from libs.datetime_utils import naive_utc_now
@ -411,7 +411,7 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline):
if trace_manager:
TelemetryFacade.emit(
TelemetryEvent(
name="message",
name=TraceTaskName.MESSAGE_TRACE,
context=TelemetryContext(
tenant_id=self._application_generate_entity.app_config.tenant_id,
app_id=self._application_generate_entity.app_config.app_id,

8
api/core/app/workflow/layers/persistence.py
View File

@ -395,11 +395,11 @@ class WorkflowPersistenceLayer(GraphEngineLayer):
external_trace_id = self._application_generate_entity.extras.get("external_trace_id")
parent_trace_context = self._application_generate_entity.extras.get("parent_trace_context")
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade, TraceTaskName
TelemetryFacade.emit(
TelemetryEvent(
name="workflow",
name=TraceTaskName.WORKFLOW_TRACE,
context=TelemetryContext(
tenant_id=self._application_generate_entity.app_config.tenant_id,
user_id=self._trace_manager.user_id,
@ -499,11 +499,11 @@ class WorkflowPersistenceLayer(GraphEngineLayer):
if parent_trace_context:
node_data["parent_trace_context"] = parent_trace_context
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade, TraceTaskName
TelemetryFacade.emit(
TelemetryEvent(
name="node_execution",
name=TraceTaskName.NODE_EXECUTION_TRACE,
context=TelemetryContext(
tenant_id=node_data.get("tenant_id"),
user_id=node_data.get("user_id"),

14
api/core/callback_handler/agent_tool_callback_handler.py
View File

@ -5,7 +5,7 @@ from pydantic import BaseModel
from configs import dify_config
from core.ops.ops_trace_manager import TraceQueueManager
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade, TraceTaskName
from core.tools.entities.tool_entities import ToolInvokeMessage
_TEXT_COLOR_MAPPING = {
@ -36,13 +36,15 @@ class DifyAgentCallbackHandler(BaseModel):
color: str | None = ""
current_loop: int = 1
tenant_id: str | None = None
def __init__(self, color: str | None = None):
def __init__(self, color: str | None = None, tenant_id: str | None = None):
super().__init__()
"""Initialize callback handler."""
# use a specific color is not specified
self.color = color or "green"
self.current_loop = 1
self.tenant_id = tenant_id
def on_tool_start(
self,
@ -73,8 +75,12 @@ class DifyAgentCallbackHandler(BaseModel):
if trace_manager:
TelemetryFacade.emit(
TelemetryEvent(
name="tool",
context=TelemetryContext(app_id=trace_manager.app_id, user_id=trace_manager.user_id),
name=TraceTaskName.TOOL_TRACE,
context=TelemetryContext(
tenant_id=self.tenant_id,
app_id=trace_manager.app_id,
user_id=trace_manager.user_id,
),
payload={
"message_id": message_id,
"tool_name": tool_name,

8
api/core/llm_generator/llm_generator.py
View File

@ -27,7 +27,7 @@ from core.model_runtime.entities.model_entities import ModelType
from core.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError
from core.ops.utils import measure_time
from core.prompt.utils.prompt_template_parser import PromptTemplateParser
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade, TraceTaskName
from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionMetadataKey
from extensions.ext_database import db
from extensions.ext_storage import storage
@ -95,7 +95,7 @@ class LLMGenerator:
# get tracing instance
TelemetryFacade.emit(
TelemetryEvent(
name="generate_name",
name=TraceTaskName.GENERATE_NAME_TRACE,
context=TelemetryContext(tenant_id=tenant_id, app_id=app_id),
payload={
"conversation_id": conversation_id,
@ -788,11 +788,9 @@ class LLMGenerator:
total_price = None
currency = None
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade
TelemetryFacade.emit(
TelemetryEvent(
name="prompt_generation",
name=TraceTaskName.PROMPT_GENERATION_TRACE,
context=TelemetryContext(tenant_id=tenant_id, user_id=user_id, app_id=app_id),
payload={
"tenant_id": tenant_id,

4
api/core/moderation/input_moderation.py
View File

@ -7,7 +7,7 @@ from core.moderation.base import ModerationAction, ModerationError
from core.moderation.factory import ModerationFactory
from core.ops.ops_trace_manager import TraceQueueManager
from core.ops.utils import measure_time
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade, TraceTaskName
logger = logging.getLogger(__name__)
@ -51,7 +51,7 @@ class InputModeration:
if trace_manager:
TelemetryFacade.emit(
TelemetryEvent(
name="moderation",
name=TraceTaskName.MODERATION_TRACE,
context=TelemetryContext(tenant_id=tenant_id, app_id=app_id),
payload={
"message_id": message_id,

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

@ -1272,9 +1272,9 @@ class TraceQueueManager:
self.trace_instance = OpsTraceManager.get_ops_trace_instance(app_id)
self.flask_app = current_app._get_current_object() # type: ignore
from core.telemetry import is_telemetry_enabled
from core.telemetry import is_enterprise_telemetry_enabled
self._enterprise_telemetry_enabled = is_telemetry_enabled()
self._enterprise_telemetry_enabled = is_enterprise_telemetry_enabled()
if trace_manager_timer is None:
self.start_timer()

4
api/core/rag/retrieval/dataset_retrieval.py
View File

@ -55,7 +55,7 @@ from core.rag.retrieval.template_prompts import (
METADATA_FILTER_USER_PROMPT_2,
METADATA_FILTER_USER_PROMPT_3,
)
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade, TraceTaskName
from core.tools.signature import sign_upload_file
from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool
from extensions.ext_database import db
@ -731,7 +731,7 @@ class DatasetRetrieval:
app_config = self.application_generate_entity.app_config if self.application_generate_entity else None
TelemetryFacade.emit(
TelemetryEvent(
name="dataset_retrieval",
name=TraceTaskName.DATASET_RETRIEVAL_TRACE,
context=TelemetryContext(
tenant_id=app_config.tenant_id if app_config else None,
app_id=app_config.app_id if app_config else None,

12
api/core/telemetry/__init__.py
View File

@ -1,4 +1,12 @@
from core.ops.entities.trace_entity import TraceTaskName
from core.telemetry.events import TelemetryContext, TelemetryEvent
from core.telemetry.facade import TelemetryFacade, emit, is_telemetry_enabled
from core.telemetry.facade import TelemetryFacade, emit, is_enterprise_telemetry_enabled
__all__ = ["TelemetryContext", "TelemetryEvent", "TelemetryFacade", "emit", "is_telemetry_enabled"]
__all__ = [
"TelemetryContext",
"TelemetryEvent",
"TelemetryFacade",
"TraceTaskName",
"emit",
"is_enterprise_telemetry_enabled",
]

7
api/core/telemetry/events.py
View File

@ -1,7 +1,10 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from core.ops.entities.trace_entity import TraceTaskName
@dataclass(frozen=True)
@ -13,6 +16,6 @@ class TelemetryContext:
@dataclass(frozen=True)
class TelemetryEvent:
name: str
name: TraceTaskName
context: TelemetryContext
payload: dict[str, Any]

31
api/core/telemetry/facade.py
View File

@ -2,32 +2,27 @@ from __future__ import annotations
from typing import TYPE_CHECKING
from core.ops.entities.trace_entity import TraceTaskName
from core.telemetry.events import TelemetryEvent
if TYPE_CHECKING:
from core.ops.ops_trace_manager import TraceQueueManager
_ENTERPRISE_ONLY_TRACES: frozenset[TraceTaskName] = frozenset(
{
TraceTaskName.DRAFT_NODE_EXECUTION_TRACE,
TraceTaskName.NODE_EXECUTION_TRACE,
TraceTaskName.PROMPT_GENERATION_TRACE,
}
)
class TelemetryFacade:
@staticmethod
def emit(event: TelemetryEvent, trace_manager: TraceQueueManager | None = None) -> None:
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask, TraceTaskName
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
trace_task_name_map = {
"draft_node_execution": TraceTaskName.DRAFT_NODE_EXECUTION_TRACE,
"dataset_retrieval": TraceTaskName.DATASET_RETRIEVAL_TRACE,
"generate_name": TraceTaskName.GENERATE_NAME_TRACE,
"message": TraceTaskName.MESSAGE_TRACE,
"moderation": TraceTaskName.MODERATION_TRACE,
"node_execution": TraceTaskName.NODE_EXECUTION_TRACE,
"prompt_generation": TraceTaskName.PROMPT_GENERATION_TRACE,
"suggested_question": TraceTaskName.SUGGESTED_QUESTION_TRACE,
"tool": TraceTaskName.TOOL_TRACE,
"workflow": TraceTaskName.WORKFLOW_TRACE,
}
trace_task_name = trace_task_name_map.get(event.name)
if not trace_task_name:
if event.name not in _ENTERPRISE_ONLY_TRACES:
return
trace_queue_manager = trace_manager or TraceQueueManager(
@ -36,13 +31,13 @@ class TelemetryFacade:
)
trace_queue_manager.add_trace_task(
TraceTask(
trace_task_name,
event.name,
**event.payload,
)
)
def is_telemetry_enabled() -> bool:
def is_enterprise_telemetry_enabled() -> bool:
try:
from enterprise.telemetry.exporter import is_enterprise_telemetry_enabled
except Exception:

8
api/enterprise/telemetry/draft_trace.py
View File

@ -3,9 +3,8 @@ from __future__ import annotations
from collections.abc import Mapping
from typing import Any
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade, TraceTaskName
from core.workflow.enums import WorkflowNodeExecutionMetadataKey
from enterprise.telemetry.exporter import is_enterprise_telemetry_enabled
from models.workflow import WorkflowNodeExecutionModel
@ -16,9 +15,6 @@ def enqueue_draft_node_execution_trace(
workflow_execution_id: str | None,
user_id: str,
) -> None:
if not is_enterprise_telemetry_enabled():
return
node_data = _build_node_execution_data(
execution=execution,
outputs=outputs,
@ -26,7 +22,7 @@ def enqueue_draft_node_execution_trace(
)
TelemetryFacade.emit(
TelemetryEvent(
name="draft_node_execution",
name=TraceTaskName.DRAFT_NODE_EXECUTION_TRACE,
context=TelemetryContext(
tenant_id=execution.tenant_id,
user_id=user_id,

4
api/services/message_service.py
View File

@ -8,7 +8,7 @@ from core.memory.token_buffer_memory import TokenBufferMemory
from core.model_manager import ModelManager
from core.model_runtime.entities.model_entities import ModelType
from core.ops.utils import measure_time
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade
from core.telemetry import TelemetryContext, TelemetryEvent, TelemetryFacade, TraceTaskName
from events.feedback_event import feedback_was_created
from extensions.ext_database import db
from libs.infinite_scroll_pagination import InfiniteScrollPagination
@ -299,7 +299,7 @@ class MessageService:
# get tracing instance
TelemetryFacade.emit(
TelemetryEvent(
name="suggested_question",
name=TraceTaskName.SUGGESTED_QUESTION_TRACE,
context=TelemetryContext(tenant_id=app_model.tenant_id, app_id=app_model.id),
payload={
"message_id": message_id,

200
api/tests/unit_tests/core/ops/test_trace_queue_manager.py Normal file
View File

@ -0,0 +1,200 @@
"""Unit tests for TraceQueueManager telemetry guard.
This test suite verifies that TraceQueueManager correctly drops trace tasks
when telemetry is disabled, proving Bug 1 from code review is a false positive.
The guard logic moved from persistence.py to TraceQueueManager.add_trace_task()
at line 1282 of ops_trace_manager.py:
if self._enterprise_telemetry_enabled or self.trace_instance:
trace_task.app_id = self.app_id
trace_manager_queue.put(trace_task)
Tasks are only enqueued if EITHER:
- Enterprise telemetry is enabled (_enterprise_telemetry_enabled=True), OR
- A third-party trace instance (Langfuse, etc.) is configured
When BOTH are false, tasks are silently dropped (correct behavior).
"""
import queue
import sys
import types
from unittest.mock import MagicMock, patch
import pytest
@pytest.fixture
def trace_queue_manager_and_task(monkeypatch):
"""Fixture to provide TraceQueueManager and TraceTask with delayed imports."""
module_name = "core.ops.ops_trace_manager"
if module_name not in sys.modules:
ops_stub = types.ModuleType(module_name)
class StubTraceTask:
def __init__(self, trace_type):
self.trace_type = trace_type
self.app_id = None
class StubTraceQueueManager:
def __init__(self, app_id=None):
self.app_id = app_id
from core.telemetry import is_enterprise_telemetry_enabled
self._enterprise_telemetry_enabled = is_enterprise_telemetry_enabled()
self.trace_instance = StubOpsTraceManager.get_ops_trace_instance(app_id)
def add_trace_task(self, trace_task):
if self._enterprise_telemetry_enabled or self.trace_instance:
trace_task.app_id = self.app_id
from core.ops.ops_trace_manager import trace_manager_queue
trace_manager_queue.put(trace_task)
class StubOpsTraceManager:
@staticmethod
def get_ops_trace_instance(app_id):
return None
ops_stub.TraceQueueManager = StubTraceQueueManager
ops_stub.TraceTask = StubTraceTask
ops_stub.OpsTraceManager = StubOpsTraceManager
ops_stub.trace_manager_queue = MagicMock(spec=queue.Queue)
monkeypatch.setitem(sys.modules, module_name, ops_stub)
from core.ops.entities.trace_entity import TraceTaskName
ops_module = __import__(module_name, fromlist=["TraceQueueManager", "TraceTask"])
TraceQueueManager = ops_module.TraceQueueManager
TraceTask = ops_module.TraceTask
return 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.
"""
TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task
mock_queue = MagicMock(spec=queue.Queue)
trace_task = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE)
with (
patch("core.telemetry.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", mock_queue),
):
manager = TraceQueueManager(app_id="test-app-id")
manager.add_trace_task(trace_task)
mock_queue.put.assert_not_called()
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.
"""
TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task
mock_queue = MagicMock(spec=queue.Queue)
trace_task = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE)
with (
patch("core.telemetry.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", mock_queue),
):
manager = TraceQueueManager(app_id="test-app-id")
manager.add_trace_task(trace_task)
mock_queue.put.assert_called_once()
called_task = mock_queue.put.call_args[0][0]
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.
"""
TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task
mock_queue = MagicMock(spec=queue.Queue)
mock_trace_instance = MagicMock()
trace_task = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE)
with (
patch("core.telemetry.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", mock_queue),
):
manager = TraceQueueManager(app_id="test-app-id")
manager.add_trace_task(trace_task)
mock_queue.put.assert_called_once()
called_task = mock_queue.put.call_args[0][0]
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.
"""
TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task
mock_queue = MagicMock(spec=queue.Queue)
mock_trace_instance = MagicMock()
trace_task = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE)
with (
patch("core.telemetry.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", mock_queue),
):
manager = TraceQueueManager(app_id="test-app-id")
manager.add_trace_task(trace_task)
mock_queue.put.assert_called_once()
called_task = mock_queue.put.call_args[0][0]
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.
"""
TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task
mock_queue = MagicMock(spec=queue.Queue)
trace_task = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE)
with (
patch("core.telemetry.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", mock_queue),
):
manager = TraceQueueManager(app_id="expected-app-id")
manager.add_trace_task(trace_task)
called_task = mock_queue.put.call_args[0][0]
assert called_task.app_id == "expected-app-id"

242
api/tests/unit_tests/core/telemetry/test_facade.py Normal file
View File

@ -0,0 +1,242 @@
"""Unit tests for TelemetryFacade.emit() routing and enterprise-only filtering.
This test suite verifies that TelemetryFacade correctly:
1. Routes telemetry events to TraceQueueManager via enum-based TraceTaskName
2. Blocks community traces (returns early)
3. Allows enterprise-only traces to be routed to TraceQueueManager
4. Passes TraceTaskName enum directly to TraceTask constructor
"""
import queue
import sys
import types
from unittest.mock import MagicMock
import pytest
from core.ops.entities.trace_entity import TraceTaskName
from core.telemetry.events import TelemetryContext, TelemetryEvent
@pytest.fixture
def facade_test_setup(monkeypatch):
"""Fixture to provide TelemetryFacade with mocked TraceQueueManager."""
module_name = "core.ops.ops_trace_manager"
# Always create a fresh stub module for testing
ops_stub = types.ModuleType(module_name)
class StubTraceTask:
def __init__(self, trace_type, **kwargs):
self.trace_type = trace_type
self.app_id = None
self.kwargs = kwargs
class StubTraceQueueManager:
def __init__(self, app_id=None, user_id=None):
self.app_id = app_id
self.user_id = user_id
self.trace_instance = StubOpsTraceManager.get_ops_trace_instance(app_id)
def add_trace_task(self, trace_task):
trace_task.app_id = self.app_id
from core.ops.ops_trace_manager import trace_manager_queue
trace_manager_queue.put(trace_task)
class StubOpsTraceManager:
@staticmethod
def get_ops_trace_instance(app_id):
return None
ops_stub.TraceQueueManager = StubTraceQueueManager
ops_stub.TraceTask = StubTraceTask
ops_stub.OpsTraceManager = StubOpsTraceManager
ops_stub.trace_manager_queue = MagicMock(spec=queue.Queue)
monkeypatch.setitem(sys.modules, module_name, ops_stub)
from core.telemetry.facade import TelemetryFacade
return TelemetryFacade, ops_stub.trace_manager_queue
class TestTelemetryFacadeEmit:
"""Test TelemetryFacade.emit() routing and filtering."""
def test_emit_valid_name_creates_trace_task(self, facade_test_setup):
"""Verify emit with enterprise-only trace creates and enqueues a trace task.
When emit() is called with an enterprise-only trace name
(DRAFT_NODE_EXECUTION_TRACE), TraceQueueManager.add_trace_task()
should be called with a TraceTask.
"""
TelemetryFacade, mock_queue = facade_test_setup
event = TelemetryEvent(
name=TraceTaskName.DRAFT_NODE_EXECUTION_TRACE,
context=TelemetryContext(
tenant_id="test-tenant",
user_id="test-user",
app_id="test-app",
),
payload={"key": "value"},
)
TelemetryFacade.emit(event)
# Verify add_trace_task was called
mock_queue.put.assert_called_once()
# Verify the TraceTask was created with the correct name
called_task = mock_queue.put.call_args[0][0]
assert called_task.trace_type == TraceTaskName.DRAFT_NODE_EXECUTION_TRACE
def test_emit_community_trace_returns_early(self, facade_test_setup):
"""Verify community trace is blocked by early return.
When emit() is called with a community trace (WORKFLOW_TRACE),
the facade should return early without calling add_trace_task.
"""
TelemetryFacade, mock_queue = facade_test_setup
event = TelemetryEvent(
name=TraceTaskName.WORKFLOW_TRACE,
context=TelemetryContext(
tenant_id="test-tenant",
user_id="test-user",
app_id="test-app",
),
payload={},
)
TelemetryFacade.emit(event)
# Community traces should not reach the queue
mock_queue.put.assert_not_called()
def test_emit_enterprise_only_trace_allowed(self, facade_test_setup):
"""Verify enterprise-only trace is routed to TraceQueueManager.
When emit() is called with DRAFT_NODE_EXECUTION_TRACE,
add_trace_task should be called.
"""
TelemetryFacade, mock_queue = facade_test_setup
event = TelemetryEvent(
name=TraceTaskName.DRAFT_NODE_EXECUTION_TRACE,
context=TelemetryContext(
tenant_id="test-tenant",
user_id="test-user",
app_id="test-app",
),
payload={},
)
TelemetryFacade.emit(event)
# Verify add_trace_task was called and task was enqueued
mock_queue.put.assert_called_once()
# Verify the TraceTask was created with the correct name
called_task = mock_queue.put.call_args[0][0]
assert called_task.trace_type == TraceTaskName.DRAFT_NODE_EXECUTION_TRACE
def test_emit_all_enterprise_only_traces_allowed(self, facade_test_setup):
"""Verify all 3 enterprise-only traces are correctly identified.
The three enterprise-only traces are:
- DRAFT_NODE_EXECUTION_TRACE
- NODE_EXECUTION_TRACE
- PROMPT_GENERATION_TRACE
When these are emitted, they should be routed to add_trace_task.
"""
TelemetryFacade, mock_queue = facade_test_setup
enterprise_only_traces = [
TraceTaskName.DRAFT_NODE_EXECUTION_TRACE,
TraceTaskName.NODE_EXECUTION_TRACE,
TraceTaskName.PROMPT_GENERATION_TRACE,
]
for trace_name in enterprise_only_traces:
mock_queue.reset_mock()
event = TelemetryEvent(
name=trace_name,
context=TelemetryContext(
tenant_id="test-tenant",
user_id="test-user",
app_id="test-app",
),
payload={},
)
TelemetryFacade.emit(event)
# All enterprise-only traces should be routed
mock_queue.put.assert_called_once()
# Verify the correct trace name was passed
called_task = mock_queue.put.call_args[0][0]
assert called_task.trace_type == trace_name
def test_emit_passes_name_directly_to_trace_task(self, facade_test_setup):
"""Verify event.name (TraceTaskName enum) is passed directly to TraceTask.
The facade should pass the TraceTaskName enum directly as the first
argument to TraceTask(), not convert it to a string.
"""
TelemetryFacade, mock_queue = facade_test_setup
event = TelemetryEvent(
name=TraceTaskName.DRAFT_NODE_EXECUTION_TRACE,
context=TelemetryContext(
tenant_id="test-tenant",
user_id="test-user",
app_id="test-app",
),
payload={"extra": "data"},
)
TelemetryFacade.emit(event)
# Verify add_trace_task was called
mock_queue.put.assert_called_once()
# Verify the TraceTask was created with the enum directly
called_task = mock_queue.put.call_args[0][0]
# The trace_type should be the enum, not a string
assert called_task.trace_type == TraceTaskName.DRAFT_NODE_EXECUTION_TRACE
assert isinstance(called_task.trace_type, TraceTaskName)
def test_emit_with_provided_trace_manager(self, facade_test_setup):
"""Verify emit uses provided trace_manager instead of creating one.
When a trace_manager is provided, emit should use it directly
instead of creating a new TraceQueueManager.
"""
TelemetryFacade, mock_queue = facade_test_setup
mock_trace_manager = MagicMock()
mock_trace_manager.add_trace_task = MagicMock()
event = TelemetryEvent(
name=TraceTaskName.NODE_EXECUTION_TRACE,
context=TelemetryContext(
tenant_id="test-tenant",
user_id="test-user",
app_id="test-app",
),
payload={},
)
TelemetryFacade.emit(event, trace_manager=mock_trace_manager)
# Verify the provided trace_manager was used
mock_trace_manager.add_trace_task.assert_called_once()
# Verify the TraceTask was created with the correct name
called_task = mock_trace_manager.add_trace_task.call_args[0][0]
assert called_task.trace_type == TraceTaskName.NODE_EXECUTION_TRACE