crewAI/lib/crewai/tests/tracing/test_tracing.py

import os
from threading import Thread
from unittest.mock import MagicMock, Mock, patch

import pytest
from crewai import Agent, Crew, Task
from crewai.events.listeners.tracing.first_time_trace_handler import (
    FirstTimeTraceHandler,
)
from crewai.events.listeners.tracing.trace_batch_manager import (
    TraceBatch,
    TraceBatchManager,
)
from crewai.events.listeners.tracing.trace_listener import (
    TraceCollectionListener,
)
from crewai.events.listeners.tracing.types import TraceEvent
from crewai.flow.flow import Flow, start
from tests.utils import wait_for_event_handlers


class TestTraceListenerSetup:
    """Test TraceListener is properly setup and collecting events"""

    @pytest.fixture(autouse=True)
    def mock_user_data_file_io(self):
        """Mock user data file I/O to prevent file system pollution between tests"""
        with patch(
            "crewai.events.listeners.tracing.utils._load_user_data",
            return_value={},
        ):
            yield

    @pytest.fixture(autouse=True)
    def mock_auth_token(self):
        """Mock authentication token for all tests in this class"""
        # Need to patch all the places where get_auth_token is imported/used
        with (
            patch(
                "crewai.auth.token.get_auth_token",
                return_value="mock_token_12345",
            ),
            patch(
                "crewai.events.listeners.tracing.trace_listener.get_auth_token",
                return_value="mock_token_12345",
            ),
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.get_auth_token",
                return_value="mock_token_12345",
            ),
        ):
            yield

    @pytest.fixture(autouse=True)
    def reset_tracing_singletons(self):
        """Reset tracing singleton instances between tests"""
        from crewai.events.event_bus import crewai_event_bus
        from crewai.events.event_listener import EventListener
        from crewai.events.listeners.tracing.utils import _tracing_enabled

        try:
            _tracing_enabled.set(None)
        except (LookupError, AttributeError):
            pass

        with crewai_event_bus._rwlock.w_locked():
            crewai_event_bus._sync_handlers = {}
            crewai_event_bus._async_handlers = {}
            crewai_event_bus._handler_dependencies = {}
            crewai_event_bus._execution_plan_cache = {}

        # Reset TraceCollectionListener singleton - must reset instance attributes too
        if TraceCollectionListener._instance is not None:
            # Reset instance attributes that shadow class attributes (only if they exist as instance attrs)
            instance_dict = TraceCollectionListener._instance.__dict__
            if "_initialized" in instance_dict:
                del TraceCollectionListener._instance._initialized
            if "_listeners_setup" in instance_dict:
                del TraceCollectionListener._instance._listeners_setup

        TraceCollectionListener._instance = None
        TraceCollectionListener._initialized = False
        TraceCollectionListener._listeners_setup = False

        if hasattr(EventListener, "_instance"):
            EventListener._instance = None

        yield

        with crewai_event_bus._rwlock.w_locked():
            crewai_event_bus._sync_handlers = {}
            crewai_event_bus._async_handlers = {}
            crewai_event_bus._handler_dependencies = {}
            crewai_event_bus._execution_plan_cache = {}

        # Reset TraceCollectionListener singleton - must reset instance attributes too
        if TraceCollectionListener._instance is not None:
            # Reset instance attributes that shadow class attributes (only if they exist as instance attrs)
            instance_dict = TraceCollectionListener._instance.__dict__
            if "_initialized" in instance_dict:
                del TraceCollectionListener._instance._initialized
            if "_listeners_setup" in instance_dict:
                del TraceCollectionListener._instance._listeners_setup

        TraceCollectionListener._instance = None
        TraceCollectionListener._initialized = False
        TraceCollectionListener._listeners_setup = False

        if hasattr(EventListener, "_instance"):
            EventListener._instance = None

    @pytest.fixture(autouse=True)
    def mock_plus_api_calls(self):
        """Mock all PlusAPI HTTP calls to avoid network requests"""
        with (
            patch("requests.post") as mock_post,
            patch("requests.get") as mock_get,
            patch("requests.put") as mock_put,
            patch("requests.delete") as mock_delete,
            patch.object(TraceBatchManager, "_cleanup_batch_data", return_value=True),
        ):
            mock_response = MagicMock()
            mock_response.status_code = 200
            mock_response.json.return_value = {
                "id": "mock_trace_batch_id",
                "status": "success",
                "message": "Batch created successfully",
            }
            mock_response.raise_for_status.return_value = None

            mock_post.return_value = mock_response
            mock_get.return_value = mock_response
            mock_put.return_value = mock_response
            mock_delete.return_value = mock_response

            mock_mark_failed = MagicMock()
            mock_mark_failed.return_value = mock_response

            yield {
                "post": mock_post,
                "get": mock_get,
                "put": mock_put,
                "delete": mock_delete,
                "mark_trace_batch_as_failed": mock_mark_failed,
            }

    @pytest.mark.vcr()
    def test_trace_listener_collects_crew_events(self):
        """Test that trace listener properly collects events from crew execution"""

        with patch.dict(
            os.environ,
            {
                "CREWAI_TRACING_ENABLED": "true",
                "CREWAI_DISABLE_TELEMETRY": "false",
                "CREWAI_DISABLE_TRACKING": "false",
                "OTEL_SDK_DISABLED": "false",
            },
        ):
            agent = Agent(
                role="Test Agent",
                goal="Test goal",
                backstory="Test backstory",
                llm="gpt-4o-mini",
            )
            task = Task(
                description="Say hello to the world",
                expected_output="hello world",
                agent=agent,
            )
            crew = Crew(agents=[agent], tasks=[task], verbose=True)

            from crewai.events.listeners.tracing.trace_listener import TraceCollectionListener
            trace_listener = TraceCollectionListener()

            crew.kickoff()

            initialized = trace_listener.batch_manager.wait_for_batch_initialization(timeout=5.0)

            assert initialized, "Batch should have been initialized"
            assert trace_listener.batch_manager.is_batch_initialized()
            assert trace_listener.batch_manager.current_batch is not None

    @pytest.mark.vcr()
    def test_batch_manager_finalizes_batch_clears_buffer(self):
        """Test that batch manager properly finalizes batch and clears buffer"""

        with patch.dict(
            os.environ,
            {
                "CREWAI_TRACING_ENABLED": "true",
                "CREWAI_DISABLE_TELEMETRY": "false",
                "CREWAI_DISABLE_TRACKING": "false",
                "OTEL_SDK_DISABLED": "false",
            },
        ):
            agent = Agent(
                role="Test Agent",
                goal="Test goal",
                backstory="Test backstory",
                llm="gpt-4o-mini",
            )

            task = Task(
                description="Say hello to the world",
                expected_output="hello world",
                agent=agent,
            )

            crew = Crew(agents=[agent], tasks=[task], verbose=True)

            from crewai.events.event_bus import crewai_event_bus

            trace_listener = None
            with crewai_event_bus._rwlock.r_locked():
                for handler_set in crewai_event_bus._sync_handlers.values():
                    for handler in handler_set:
                        if hasattr(handler, "__self__") and isinstance(
                            handler.__self__, TraceCollectionListener
                        ):
                            trace_listener = handler.__self__
                            break
                    if trace_listener:
                        break
                if not trace_listener:
                    for handler_set in crewai_event_bus._async_handlers.values():
                        for handler in handler_set:
                            if hasattr(handler, "__self__") and isinstance(
                                handler.__self__, TraceCollectionListener
                            ):
                                trace_listener = handler.__self__
                                break
                        if trace_listener:
                            break

            if not trace_listener:
                pytest.skip(
                    "No trace listener found - tracing may not be properly enabled"
                )

            with patch.object(
                trace_listener.batch_manager,
                "finalize_batch",
                wraps=trace_listener.batch_manager.finalize_batch,
            ) as finalize_mock:
                crew.kickoff()

                assert finalize_mock.call_count >= 1

    @pytest.mark.vcr()
    def test_events_collection_batch_manager(self, mock_plus_api_calls):
        """Test that trace listener properly collects events from crew execution"""

        with patch.dict(
            os.environ,
            {
                "CREWAI_TRACING_ENABLED": "true",
                "CREWAI_DISABLE_TELEMETRY": "false",
                "CREWAI_DISABLE_TRACKING": "false",
                "OTEL_SDK_DISABLED": "false",
            },
        ):
            agent = Agent(
                role="Test Agent",
                goal="Test goal",
                backstory="Test backstory",
                llm="gpt-4o-mini",
            )
            task = Task(
                description="Say hello to the world",
                expected_output="hello world",
                agent=agent,
            )
            crew = Crew(agents=[agent], tasks=[task], verbose=True)

            from crewai.events.event_bus import crewai_event_bus

            trace_listener = TraceCollectionListener()
            trace_listener.setup_listeners(crewai_event_bus)

            with patch.object(
                trace_listener.batch_manager,
                "add_event",
                wraps=trace_listener.batch_manager.add_event,
            ) as add_event_mock:
                crew.kickoff()
                wait_for_event_handlers()

                assert add_event_mock.call_count >= 2

                completion_events = [
                    call.args[0]
                    for call in add_event_mock.call_args_list
                    if call.args[0].type == "crew_kickoff_completed"
                ]
                assert len(completion_events) >= 1

                completion_event = completion_events[0]
                assert "crew_name" in completion_event.event_data
                assert completion_event.event_data["crew_name"] == "crew"

                for call in add_event_mock.call_args_list:
                    event = call.args[0]
                    assert isinstance(event, TraceEvent)
                    assert hasattr(event, "event_data")
                    assert hasattr(event, "type")

    @pytest.mark.vcr()
    def test_trace_listener_disabled_when_env_false(self):
        """Test that trace listener doesn't make HTTP calls when tracing is disabled"""

        with patch.dict(os.environ, {"CREWAI_TRACING_ENABLED": "false"}):
            agent = Agent(
                role="Test Agent",
                goal="Test goal",
                backstory="Test backstory",
                llm="gpt-4o-mini",
            )
            task = Task(
                description="Say hello to the world",
                expected_output="hello world",
                agent=agent,
            )

            crew = Crew(agents=[agent], tasks=[task], verbose=True)
            result = crew.kickoff()
            assert result is not None

            from crewai.events.event_bus import crewai_event_bus

            trace_handlers = []
            with crewai_event_bus._rwlock.r_locked():
                for handlers in crewai_event_bus._sync_handlers.values():
                    for handler in handlers:
                        if hasattr(handler, "__self__") and isinstance(
                            handler.__self__, TraceCollectionListener
                        ):
                            trace_handlers.append(handler)
                for handlers in crewai_event_bus._async_handlers.values():
                    for handler in handlers:
                        if hasattr(handler, "__self__") and isinstance(
                            handler.__self__, TraceCollectionListener
                        ):
                            trace_handlers.append(handler)

            assert len(trace_handlers) == 0, (
                f"Found {len(trace_handlers)} TraceCollectionListener handlers when tracing should be disabled"
            )

    def test_trace_listener_setup_correctly_for_crew(self):
        """Test that trace listener is set up correctly when enabled"""

        with patch.dict(
            os.environ,
            {
                "CREWAI_TRACING_ENABLED": "true",
                "CREWAI_DISABLE_TELEMETRY": "false",
                "CREWAI_DISABLE_TRACKING": "false",
                "OTEL_SDK_DISABLED": "false",
            },
        ):
            agent = Agent(
                role="Test Agent",
                goal="Test goal",
                backstory="Test backstory",
                llm="gpt-4o-mini",
            )
            task = Task(
                description="Say hello to the world",
                expected_output="hello world",
                agent=agent,
            )
            with patch.object(
                TraceCollectionListener, "setup_listeners"
            ) as mock_listener_setup:
                Crew(agents=[agent], tasks=[task], verbose=True)
                assert mock_listener_setup.call_count >= 1

    @pytest.mark.vcr()
    def test_trace_listener_setup_correctly_for_flow(self):
        """Test that trace listener is set up correctly when enabled"""

        with patch.dict(
            os.environ,
            {
                "CREWAI_TRACING_ENABLED": "true",
                "CREWAI_DISABLE_TELEMETRY": "false",
                "CREWAI_DISABLE_TRACKING": "false",
                "OTEL_SDK_DISABLED": "false",
            },
        ):
            class FlowExample(Flow):
                @start()
                def start(self):
                    pass

            with patch.object(
                TraceCollectionListener, "setup_listeners"
            ) as mock_listener_setup:
                FlowExample()
                assert mock_listener_setup.call_count >= 1

    @pytest.mark.vcr()
    def test_trace_listener_ephemeral_batch(self):
        """Test that trace listener properly handles ephemeral batches"""
        with (
            patch.dict(
                os.environ,
                {
                    "CREWAI_TRACING_ENABLED": "true",
                    "CREWAI_DISABLE_TELEMETRY": "false",
                    "CREWAI_DISABLE_TRACKING": "false",
                    "OTEL_SDK_DISABLED": "false",
                },
            ),
            patch(
                "crewai.events.listeners.tracing.trace_listener.TraceCollectionListener._check_authenticated",
                return_value=False,
            ),
        ):
            agent = Agent(
                role="Test Agent",
                goal="Test goal",
                backstory="Test backstory",
                llm="gpt-4o-mini",
            )
            task = Task(
                description="Say hello to the world",
                expected_output="hello world",
                agent=agent,
            )
            crew = Crew(agents=[agent], tasks=[task], tracing=True)

            from crewai.events.listeners.tracing.trace_listener import TraceCollectionListener
            trace_listener = TraceCollectionListener()

            crew.kickoff()

            initialized = trace_listener.batch_manager.wait_for_batch_initialization(timeout=5.0)
            assert initialized, (
                "Batch should have been initialized for unauthenticated user"
            )

            wait_for_event_handlers()

    @pytest.mark.vcr()
    def test_trace_listener_with_authenticated_user(self):
        """Test that trace listener properly handles authenticated batches"""
        with patch.dict(
            os.environ,
            {
                "CREWAI_TRACING_ENABLED": "true",
                "CREWAI_DISABLE_TELEMETRY": "false",
                "CREWAI_DISABLE_TRACKING": "false",
                "OTEL_SDK_DISABLED": "false",
            },
        ):
            agent = Agent(
                role="Test Agent",
                goal="Test goal",
                backstory="Test backstory",
                llm="gpt-4o-mini",
            )
            task = Task(
                description="Say hello to the world",
                expected_output="hello world",
                agent=agent,
            )

            from crewai.events.listeners.tracing.trace_listener import TraceCollectionListener
            trace_listener = TraceCollectionListener()

            crew = Crew(agents=[agent], tasks=[task], tracing=True)
            crew.kickoff()

            initialized = trace_listener.batch_manager.wait_for_batch_initialization(timeout=5.0)
            assert initialized, (
                "Batch should have been initialized for authenticated user"
            )

            wait_for_event_handlers()

    # Helper method to ensure cleanup
    def teardown_method(self):
        """Cleanup after each test method"""
        from crewai.events.event_bus import crewai_event_bus
        from crewai.events.event_listener import EventListener

        with crewai_event_bus._rwlock.w_locked():
            crewai_event_bus._sync_handlers = {}
            crewai_event_bus._async_handlers = {}
            crewai_event_bus._handler_dependencies = {}
            crewai_event_bus._execution_plan_cache = {}

        if hasattr(EventListener, "_instance"):
            EventListener._instance = None

    @classmethod
    def teardown_class(cls):
        """Final cleanup after all tests in this class"""
        from crewai.events.event_bus import crewai_event_bus
        from crewai.events.event_listener import EventListener

        with crewai_event_bus._rwlock.w_locked():
            crewai_event_bus._sync_handlers = {}
            crewai_event_bus._async_handlers = {}
            crewai_event_bus._handler_dependencies = {}
            crewai_event_bus._execution_plan_cache = {}

        if hasattr(EventListener, "_instance"):
            EventListener._instance = None

    @pytest.mark.vcr()
    def test_first_time_user_trace_collection_with_timeout(self, mock_plus_api_calls):
        """Test first-time user trace collection logic with timeout behavior"""

        with (
            patch.dict(
                os.environ,
                {
                    "CREWAI_TRACING_ENABLED": "false",
                    "CREWAI_DISABLE_TELEMETRY": "false",
                    "CREWAI_DISABLE_TRACKING": "false",
                    "OTEL_SDK_DISABLED": "false",
                },
            ),
            patch(
                "crewai.events.listeners.tracing.utils._is_test_environment",
                return_value=False,
            ),
            patch(
                "crewai.events.listeners.tracing.utils.should_auto_collect_first_time_traces",
                return_value=True,
            ),
            patch(
                "crewai.events.listeners.tracing.utils.is_first_execution",
                return_value=True,
            ),
            patch(
                "crewai.events.listeners.tracing.first_time_trace_handler.prompt_user_for_trace_viewing",
                return_value=False,
            ) as mock_prompt,
            patch(
                "crewai.events.listeners.tracing.first_time_trace_handler.mark_first_execution_completed"
            ) as mock_mark_completed,
        ):
            agent = Agent(
                role="Test Agent",
                goal="Test goal",
                backstory="Test backstory",
                llm="gpt-4o-mini",
            )
            task = Task(
                description="Say hello to the world",
                expected_output="hello world",
                agent=agent,
            )
            crew = Crew(agents=[agent], tasks=[task], verbose=True)

            from crewai.events.event_bus import crewai_event_bus

            trace_listener = TraceCollectionListener()
            trace_listener.setup_listeners(crewai_event_bus)

            trace_listener.first_time_handler = FirstTimeTraceHandler()
            if trace_listener.first_time_handler.initialize_for_first_time_user():
                trace_listener.first_time_handler.set_batch_manager(trace_listener.batch_manager)

            assert trace_listener.first_time_handler.is_first_time is True
            assert trace_listener.first_time_handler.collected_events is False

            trace_listener.batch_manager.batch_owner_type = "crew"

            result = crew.kickoff()
            wait_for_event_handlers()
            assert result is not None

            assert trace_listener.first_time_handler.collected_events is True, (
                "Events should have been collected"
            )

            mock_prompt.assert_called_once()

            mock_mark_completed.assert_called_once()

    @pytest.mark.vcr()
    def test_first_time_user_trace_collection_user_accepts(self, mock_plus_api_calls):
        """Test first-time user trace collection when user accepts viewing traces"""

        with (
            patch.dict(
                os.environ,
                {
                    "CREWAI_TRACING_ENABLED": "false",
                    "CREWAI_DISABLE_TELEMETRY": "false",
                    "CREWAI_DISABLE_TRACKING": "false",
                    "OTEL_SDK_DISABLED": "false",
                },
            ),
            patch(
                "crewai.events.listeners.tracing.utils._is_test_environment",
                return_value=False,
            ),
            patch(
                "crewai.events.listeners.tracing.utils.should_auto_collect_first_time_traces",
                return_value=True,
            ),
            patch(
                "crewai.events.listeners.tracing.utils.is_first_execution",
                return_value=True,
            ),
            patch(
                "crewai.events.listeners.tracing.first_time_trace_handler.prompt_user_for_trace_viewing",
                return_value=True,
            ),
            patch(
                "crewai.events.listeners.tracing.first_time_trace_handler.mark_first_execution_completed"
            ) as mock_mark_completed,
        ):
            agent = Agent(
                role="Test Agent",
                goal="Test goal",
                backstory="Test backstory",
                llm="gpt-4o-mini",
            )
            task = Task(
                description="Say hello to the world",
                expected_output="hello world",
                agent=agent,
            )
            crew = Crew(agents=[agent], tasks=[task], verbose=True)

            from crewai.events.event_bus import crewai_event_bus

            trace_listener = TraceCollectionListener()
            trace_listener.setup_listeners(crewai_event_bus)

            # Re-initialize first-time handler after patches are applied to ensure clean state
            trace_listener.first_time_handler = FirstTimeTraceHandler()
            if trace_listener.first_time_handler.initialize_for_first_time_user():
                trace_listener.first_time_handler.set_batch_manager(trace_listener.batch_manager)

            trace_listener.batch_manager.ephemeral_trace_url = (
                "https://crewai.com/trace/mock-id"
            )

            assert trace_listener.first_time_handler.is_first_time is True

            trace_listener.first_time_handler.collected_events = True

            mock_batch_response = MagicMock()
            mock_batch_response.status_code = 201
            mock_batch_response.json.return_value = {
                "trace_id": "mock-trace-id",
                "ephemeral_trace_id": "mock-ephemeral-trace-id",
                "access_code": "TRACE-mock",
            }
            mock_events_response = MagicMock()
            mock_events_response.status_code = 200

            with (
                patch.object(
                    trace_listener.first_time_handler,
                    "_initialize_backend_and_send_events",
                    wraps=trace_listener.first_time_handler._initialize_backend_and_send_events,
                ) as mock_init_backend,
                patch.object(
                    trace_listener.first_time_handler, "_display_ephemeral_trace_link"
                ) as mock_display_link,
                patch.object(
                    trace_listener.batch_manager.plus_api,
                    "initialize_trace_batch",
                    return_value=mock_batch_response,
                ),
                patch.object(
                    trace_listener.batch_manager.plus_api,
                    "initialize_ephemeral_trace_batch",
                    return_value=mock_batch_response,
                ),
                patch.object(
                    trace_listener.batch_manager.plus_api,
                    "send_trace_events",
                    return_value=mock_events_response,
                ),
                patch.object(
                    trace_listener.batch_manager.plus_api,
                    "send_ephemeral_trace_events",
                    return_value=mock_events_response,
                ),
                patch.object(
                    trace_listener.batch_manager.plus_api,
                    "finalize_trace_batch",
                    return_value=mock_events_response,
                ),
                patch.object(
                    trace_listener.batch_manager.plus_api,
                    "finalize_ephemeral_trace_batch",
                    return_value=mock_events_response,
                ),
                patch.object(
                    trace_listener.batch_manager,
                    "_cleanup_batch_data",
                ),
            ):
                crew.kickoff()
                wait_for_event_handlers()

                mock_init_backend.assert_called_once()

                mock_display_link.assert_called_once()

            mock_mark_completed.assert_called_once()

    @pytest.mark.vcr()
    def test_first_time_user_trace_consolidation_logic(self, mock_plus_api_calls):
        """Test the consolidation logic for first-time users vs regular tracing"""
        with (
            patch.dict(
                os.environ,
                {
                    "CREWAI_TRACING_ENABLED": "",
                    "CREWAI_DISABLE_TELEMETRY": "false",
                    "CREWAI_DISABLE_TRACKING": "false",
                    "OTEL_SDK_DISABLED": "false",
                },
            ),
            patch(
                "crewai.events.listeners.tracing.utils._is_test_environment",
                return_value=False,
            ),
            patch(
                "crewai.events.listeners.tracing.utils.should_auto_collect_first_time_traces",
                return_value=True,
            ),
            patch(
                "crewai.events.listeners.tracing.utils.is_first_execution",
                return_value=True,
            ),
        ):
            from crewai.events.event_bus import crewai_event_bus

            with crewai_event_bus._rwlock.w_locked():
                crewai_event_bus._sync_handlers = {}
                crewai_event_bus._async_handlers = {}

            trace_listener = TraceCollectionListener()

            # Re-initialize first-time handler after patches are applied to ensure clean state
            # This is necessary because the singleton may have been created before patches were active
            trace_listener.first_time_handler = FirstTimeTraceHandler()
            if trace_listener.first_time_handler.initialize_for_first_time_user():
                trace_listener.first_time_handler.set_batch_manager(trace_listener.batch_manager)

            trace_listener.setup_listeners(crewai_event_bus)

            assert trace_listener.first_time_handler.is_first_time is True

            agent = Agent(
                role="Test Agent",
                goal="Test goal",
                backstory="Test backstory",
                llm="gpt-4o-mini",
            )
            task = Task(
                description="Test task", expected_output="test output", agent=agent
            )
            crew = Crew(agents=[agent], tasks=[task])

            result = crew.kickoff()

            wait_for_event_handlers()

            assert trace_listener.batch_manager.is_batch_initialized(), (
                "Batch should have been initialized for first-time user"
            )
            assert result is not None

    def test_first_time_handler_timeout_behavior(self):
        """Test the timeout behavior of the first-time trace prompt"""

        with (
            patch(
                "crewai.events.listeners.tracing.utils._is_test_environment",
                return_value=False,
            ),
            patch(
                "crewai.events.listeners.tracing.utils._is_interactive_terminal",
                return_value=True,
            ),
            patch("threading.Thread") as mock_thread,
        ):
            from crewai.events.listeners.tracing.utils import (
                prompt_user_for_trace_viewing,
            )

            mock_thread_instance = Mock()
            mock_thread_instance.is_alive.return_value = True
            mock_thread.return_value = mock_thread_instance

            result = prompt_user_for_trace_viewing(timeout_seconds=5)

            assert result is False
            mock_thread.assert_called_once()
            call_args = mock_thread.call_args
            assert call_args[1]["daemon"] is True

            mock_thread_instance.start.assert_called_once()
            mock_thread_instance.join.assert_called_once_with(timeout=5)
            mock_thread_instance.is_alive.assert_called_once()

    def test_first_time_handler_graceful_error_handling(self):
        """Test graceful error handling in first-time trace logic"""

        with (
            patch(
                "crewai.events.listeners.tracing.utils.should_auto_collect_first_time_traces",
                return_value=True,
            ),
            patch(
                "crewai.events.listeners.tracing.first_time_trace_handler.prompt_user_for_trace_viewing",
                side_effect=Exception("Prompt failed"),
            ),
            patch(
                "crewai.events.listeners.tracing.first_time_trace_handler.mark_first_execution_completed"
            ) as mock_mark_completed,
        ):
            handler = FirstTimeTraceHandler()
            handler.is_first_time = True
            handler.collected_events = True

            handler.handle_execution_completion()

            mock_mark_completed.assert_called_once()

    def test_trace_batch_marked_as_failed_on_finalize_error(self):
        """Test that trace batch is marked as failed when finalization returns non-200 status"""
        with patch("crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context", return_value=True):
            batch_manager = TraceBatchManager()

            batch_manager.current_batch = batch_manager.initialize_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={
                    "execution_type": "crew",
                    "crew_name": "test_crew",
                },
            )
            batch_manager.trace_batch_id = "test_batch_id_12345"
            batch_manager.backend_initialized = True

            with (
                patch.object(
                    batch_manager.plus_api,
                    "send_trace_events",
                    return_value=MagicMock(status_code=200),
                ),
                patch.object(
                    batch_manager.plus_api,
                    "finalize_trace_batch",
                    return_value=MagicMock(status_code=500, text="Internal Server Error"),
                ),
                patch.object(
                    batch_manager.plus_api,
                    "mark_trace_batch_as_failed",
                ) as mock_mark_failed,
            ):
                batch_manager.finalize_batch()

                mock_mark_failed.assert_called_once_with(
                    "test_batch_id_12345", "Internal Server Error"
                )
                assert batch_manager.current_batch is not None
                assert batch_manager.trace_batch_id == "test_batch_id_12345"
                assert batch_manager._batch_finalized is False

    def test_finalize_batch_clears_buffer_after_successful_send(self) -> None:
        """Successful send must not restore a stale event buffer (duplicate events)."""
        from crewai.events.listeners.tracing.types import TraceEvent

        with patch(
            "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
            return_value=True,
        ):
            batch_manager = TraceBatchManager()
            batch_manager.current_batch = batch_manager.initialize_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={
                    "execution_type": "flow",
                    "flow_name": "TestFlow",
                },
            )
            batch_manager.trace_batch_id = "batch-clear-test"
            batch_manager.backend_initialized = True
            batch_manager.event_buffer = [
                TraceEvent(
                    type="llm_call_started",
                    timestamp="2026-01-01T00:00:00",
                    event_id="evt-1",
                    emission_sequence=1,
                )
            ]

            with (
                patch.object(
                    batch_manager.plus_api,
                    "send_trace_events",
                    return_value=MagicMock(status_code=200),
                ),
                patch.object(
                    batch_manager.plus_api,
                    "finalize_trace_batch",
                    return_value=MagicMock(status_code=200, json=MagicMock(return_value={})),
                ),
            ):
                batch_manager.finalize_batch()

            assert batch_manager.event_buffer == []

    def test_finalize_backend_batch_uses_captured_batch_id_for_ephemeral_panel(
        self,
    ) -> None:
        """Finalization output must not render None if manager state is reset."""
        batch_manager = TraceBatchManager()
        batch_manager.trace_batch_id = "ephemeral-batch-id"
        batch_manager.is_current_batch_ephemeral = True

        def clear_batch_id_during_response() -> dict[str, str]:
            batch_manager.trace_batch_id = None
            return {"access_code": "TRACE-test"}

        with (
            patch.object(
                batch_manager.plus_api,
                "finalize_ephemeral_trace_batch",
                return_value=MagicMock(
                    status_code=200,
                    json=clear_batch_id_during_response,
                ),
            ),
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.should_auto_collect_first_time_traces",
                return_value=False,
            ),
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.Console.print"
            ) as mock_print,
        ):
            assert batch_manager._finalize_backend_batch() is True

        panel = mock_print.call_args.args[0]
        panel_text = str(panel.renderable)
        assert "session ID: ephemeral-batch-id" in panel_text
        assert "ephemeral_trace_batches/ephemeral-batch-id" in panel_text
        assert "session ID: None" not in panel_text
        assert "ephemeral_trace_batches/None" not in panel_text

    def test_finalize_backend_batch_is_serialized(self) -> None:
        """Concurrent finalizers must only call the backend once."""
        batch_manager = TraceBatchManager()
        batch_manager.trace_batch_id = "ephemeral-batch-id"
        batch_manager.is_current_batch_ephemeral = True
        response = MagicMock(status_code=200, json=MagicMock(return_value={}))

        with (
            patch.object(
                batch_manager.plus_api,
                "finalize_ephemeral_trace_batch",
                return_value=response,
            ) as mock_finalize,
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.should_auto_collect_first_time_traces",
                return_value=True,
            ),
        ):
            results: list[bool] = []

            def finalize() -> None:
                results.append(batch_manager._finalize_backend_batch())

            threads = [Thread(target=finalize), Thread(target=finalize)]
            for thread in threads:
                thread.start()
            for thread in threads:
                thread.join()

        assert results == [True, True]
        mock_finalize.assert_called_once()

    def test_ephemeral_batch_includes_anon_id(self):
        """Test that ephemeral batch initialization sends anon_id from get_user_id()"""
        fake_user_id = "abc123def456"

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.get_user_id",
                return_value=fake_user_id,
            ),
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.should_auto_collect_first_time_traces",
                return_value=False,
            ),
        ):
            batch_manager = TraceBatchManager()

            mock_response = MagicMock(
                status_code=201,
                json=MagicMock(return_value={
                    "ephemeral_trace_id": "test-trace-id",
                    "access_code": "TRACE-abc123",
                }),
            )

            with patch.object(
                batch_manager.plus_api,
                "initialize_ephemeral_trace_batch",
                return_value=mock_response,
            ) as mock_init:
                batch_manager.initialize_batch(
                    user_context={"privacy_level": "standard"},
                    execution_metadata={
                        "execution_type": "crew",
                        "crew_name": "test_crew",
                    },
                    use_ephemeral=True,
                )

                mock_init.assert_called_once()
                payload = mock_init.call_args[0][0]
                assert payload["user_identifier"] == fake_user_id
                assert "ephemeral_trace_id" in payload

    def test_non_ephemeral_batch_does_not_include_anon_id(self):
        """Test that non-ephemeral batch initialization does not send anon_id"""
        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.should_auto_collect_first_time_traces",
                return_value=False,
            ),
        ):
            batch_manager = TraceBatchManager()

            mock_response = MagicMock(
                status_code=201,
                json=MagicMock(return_value={"trace_id": "test-trace-id"}),
            )

            with patch.object(
                batch_manager.plus_api,
                "initialize_trace_batch",
                return_value=mock_response,
            ) as mock_init:
                batch_manager.initialize_batch(
                    user_context={"privacy_level": "standard"},
                    execution_metadata={
                        "execution_type": "crew",
                        "crew_name": "test_crew",
                    },
                    use_ephemeral=False,
                )

                mock_init.assert_called_once()
                payload = mock_init.call_args[0][0]
                assert "user_identifier" not in payload


class TestTraceBatchIdClearedOnFailure:
    """Tests: trace_batch_id is cleared when _initialize_backend_batch fails."""

    def _make_batch_manager(self):
        """Create a TraceBatchManager with a pre-set trace_batch_id (simulating first-time user)."""
        with patch(
            "crewai.events.listeners.tracing.trace_batch_manager.get_auth_token",
            return_value="mock_token",
        ):
            bm = TraceBatchManager()
        bm.current_batch = TraceBatch(
            user_context={"privacy_level": "standard"},
            execution_metadata={"execution_type": "crew", "crew_name": "test"},
        )
        bm.trace_batch_id = bm.current_batch.batch_id  # simulate line 96
        bm.is_current_batch_ephemeral = True
        return bm

    def test_trace_batch_id_cleared_on_exception(self):
        """trace_batch_id must be None when the API call raises an exception."""
        bm = self._make_batch_manager()
        assert bm.trace_batch_id is not None

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                side_effect=ConnectionError("network down"),
            ),
        ):
            bm._initialize_backend_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={"execution_type": "crew"},
                use_ephemeral=True,
            )

        assert bm.trace_batch_id is None

    def test_trace_batch_id_set_on_success(self):
        """trace_batch_id must be set from the server response on success."""
        bm = self._make_batch_manager()
        server_id = "server-ephemeral-trace-id-999"

        mock_response = MagicMock(
            status_code=201,
            json=MagicMock(return_value={"ephemeral_trace_id": server_id}),
        )

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                return_value=mock_response,
            ),
        ):
            bm._initialize_backend_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={"execution_type": "crew"},
                use_ephemeral=True,
            )

        assert bm.trace_batch_id == server_id

    def test_send_events_skipped_when_trace_batch_id_none(self):
        """_send_events_to_backend must return early when trace_batch_id is None."""
        bm = self._make_batch_manager()
        bm.trace_batch_id = None
        bm.event_buffer = [MagicMock()]

        with patch.object(
            bm.plus_api, "send_ephemeral_trace_events"
        ) as mock_send:
            result = bm._send_events_to_backend()

        assert result == 500
        mock_send.assert_not_called()


class TestInitializeBackendBatchRetry:
    """Tests for retry logic in _initialize_backend_batch."""

    def _make_batch_manager(self):
        """Create a TraceBatchManager with a pre-set trace_batch_id."""
        with patch(
            "crewai.events.listeners.tracing.trace_batch_manager.get_auth_token",
            return_value="mock_token",
        ):
            bm = TraceBatchManager()
        bm.current_batch = TraceBatch(
            user_context={"privacy_level": "standard"},
            execution_metadata={"execution_type": "crew", "crew_name": "test"},
        )
        bm.trace_batch_id = bm.current_batch.batch_id
        bm.is_current_batch_ephemeral = True
        return bm

    def test_retries_on_none_response_then_succeeds(self):
        """Retries when API returns None, succeeds on second attempt."""
        bm = self._make_batch_manager()
        server_id = "server-id-after-retry"

        success_response = MagicMock(
            status_code=201,
            json=MagicMock(return_value={"ephemeral_trace_id": server_id}),
        )

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                side_effect=[None, success_response],
            ) as mock_init,
            patch("crewai.events.listeners.tracing.trace_batch_manager.time.sleep") as mock_sleep,
        ):
            bm._initialize_backend_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={"execution_type": "crew"},
                use_ephemeral=True,
            )

        assert bm.trace_batch_id == server_id
        assert mock_init.call_count == 2
        mock_sleep.assert_called_once_with(0.2)

    def test_retries_on_5xx_then_succeeds(self):
        """Retries on 500 server error, succeeds on second attempt."""
        bm = self._make_batch_manager()
        server_id = "server-id-after-5xx"

        error_response = MagicMock(status_code=500, text="Internal Server Error")
        success_response = MagicMock(
            status_code=201,
            json=MagicMock(return_value={"ephemeral_trace_id": server_id}),
        )

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                side_effect=[error_response, success_response],
            ) as mock_init,
            patch("crewai.events.listeners.tracing.trace_batch_manager.time.sleep"),
        ):
            bm._initialize_backend_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={"execution_type": "crew"},
                use_ephemeral=True,
            )

        assert bm.trace_batch_id == server_id
        assert mock_init.call_count == 2

    def test_no_retry_on_exception(self):
        """Exceptions (e.g. timeout, connection error) abort immediately without retry."""
        bm = self._make_batch_manager()

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                side_effect=ConnectionError("network down"),
            ) as mock_init,
            patch("crewai.events.listeners.tracing.trace_batch_manager.time.sleep") as mock_sleep,
        ):
            bm._initialize_backend_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={"execution_type": "crew"},
                use_ephemeral=True,
            )

        assert bm.trace_batch_id is None
        assert mock_init.call_count == 1
        mock_sleep.assert_not_called()

    def test_no_retry_on_4xx(self):
        """Does NOT retry on 422 — client error is not transient."""
        bm = self._make_batch_manager()

        error_response = MagicMock(status_code=422, text="Unprocessable Entity")

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                return_value=error_response,
            ) as mock_init,
            patch("crewai.events.listeners.tracing.trace_batch_manager.time.sleep") as mock_sleep,
        ):
            bm._initialize_backend_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={"execution_type": "crew"},
                use_ephemeral=True,
            )

        assert bm.trace_batch_id is None
        assert mock_init.call_count == 1
        mock_sleep.assert_not_called()

    def test_exhausts_retries_then_clears_batch_id(self):
        """After all retries fail, trace_batch_id is None."""
        bm = self._make_batch_manager()

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                return_value=None,
            ) as mock_init,
            patch("crewai.events.listeners.tracing.trace_batch_manager.time.sleep"),
        ):
            bm._initialize_backend_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={"execution_type": "crew"},
                use_ephemeral=True,
            )

        assert bm.trace_batch_id is None
        assert mock_init.call_count == 2  # initial + 1 retry


class TestFirstTimeHandlerBackendInitGuard:
    """Tests: backend_initialized gated on actual batch creation success."""

    def _make_handler_with_manager(self):
        """Create a FirstTimeTraceHandler wired to a TraceBatchManager."""
        with patch(
            "crewai.events.listeners.tracing.trace_batch_manager.get_auth_token",
            return_value="mock_token",
        ):
            bm = TraceBatchManager()
        bm.current_batch = TraceBatch(
            user_context={"privacy_level": "standard"},
            execution_metadata={"execution_type": "crew", "crew_name": "test"},
        )
        bm.trace_batch_id = bm.current_batch.batch_id
        bm.is_current_batch_ephemeral = True

        handler = FirstTimeTraceHandler()
        handler.is_first_time = True
        handler.collected_events = True
        handler.batch_manager = bm
        return handler, bm

    def test_backend_initialized_true_on_success(self):
        """Events are sent when batch creation succeeds, then state is cleaned up."""
        handler, bm = self._make_handler_with_manager()
        server_id = "server-id-abc"

        mock_init_response = MagicMock(
            status_code=201,
            json=MagicMock(return_value={"ephemeral_trace_id": server_id}),
        )
        mock_send_response = MagicMock(status_code=200)

        trace_batch_id_during_send = None

        def capture_send(*args, **kwargs):
            nonlocal trace_batch_id_during_send
            trace_batch_id_during_send = bm.trace_batch_id
            return mock_send_response

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                return_value=mock_init_response,
            ),
            patch.object(
                bm.plus_api,
                "send_ephemeral_trace_events",
                side_effect=capture_send,
            ),
            patch.object(bm, "_finalize_backend_batch"),
        ):
            bm.event_buffer = [MagicMock(to_dict=MagicMock(return_value={}))]
            handler._initialize_backend_and_send_events()

        # trace_batch_id was set correctly during send
        assert trace_batch_id_during_send == server_id
        # State cleaned up after completion (singleton reuse)
        assert bm.backend_initialized is False
        assert bm.trace_batch_id is None
        assert bm.current_batch is None

    def test_backend_initialized_false_on_failure(self):
        """backend_initialized stays False and events are NOT sent when batch creation fails."""
        handler, bm = self._make_handler_with_manager()

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                return_value=None,
            ),
            patch.object(bm, "_send_events_to_backend") as mock_send,
            patch.object(bm, "_finalize_backend_batch") as mock_finalize,
            patch.object(handler, "_gracefully_fail") as mock_fail,
        ):
            bm.event_buffer = [MagicMock()]
            handler._initialize_backend_and_send_events()

        assert bm.backend_initialized is False
        assert bm.trace_batch_id is None
        mock_send.assert_not_called()
        mock_finalize.assert_not_called()
        mock_fail.assert_called_once()

    def test_backend_initialized_false_on_non_2xx(self):
        """backend_initialized stays False when server returns non-2xx."""
        handler, bm = self._make_handler_with_manager()

        mock_response = MagicMock(status_code=500, text="Internal Server Error")

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                return_value=mock_response,
            ),
            patch.object(bm, "_send_events_to_backend") as mock_send,
            patch.object(bm, "_finalize_backend_batch") as mock_finalize,
            patch.object(handler, "_gracefully_fail") as mock_fail,
        ):
            bm.event_buffer = [MagicMock()]
            handler._initialize_backend_and_send_events()

        assert bm.backend_initialized is False
        assert bm.trace_batch_id is None
        mock_send.assert_not_called()
        mock_finalize.assert_not_called()
        mock_fail.assert_called_once()


class TestFirstTimeHandlerAlwaysEphemeral:
    """Tests that first-time handler always uses ephemeral with skip_context_check."""

    def _make_handler_with_manager(self):
        with patch(
            "crewai.events.listeners.tracing.trace_batch_manager.get_auth_token",
            return_value="mock_token",
        ):
            bm = TraceBatchManager()
        bm.current_batch = TraceBatch(
            user_context={"privacy_level": "standard"},
            execution_metadata={"execution_type": "crew", "crew_name": "test"},
        )
        bm.trace_batch_id = bm.current_batch.batch_id
        bm.is_current_batch_ephemeral = True

        handler = FirstTimeTraceHandler()
        handler.is_first_time = True
        handler.collected_events = True
        handler.batch_manager = bm
        return handler, bm

    def test_deferred_init_uses_ephemeral_and_skip_context_check(self):
        """Deferred backend init always uses ephemeral=True and skip_context_check=True."""
        handler, bm = self._make_handler_with_manager()

        with (
            patch.object(bm, "_initialize_backend_batch") as mock_init,
            patch.object(bm, "_send_events_to_backend"),
            patch.object(bm, "_finalize_backend_batch"),
        ):
            mock_init.side_effect = lambda **kwargs: None
            bm.event_buffer = [MagicMock()]
            handler._initialize_backend_and_send_events()

            mock_init.assert_called_once()
            assert mock_init.call_args.kwargs["use_ephemeral"] is True
            assert mock_init.call_args.kwargs["skip_context_check"] is True


class TestAuthFailbackToEphemeral:
    """Tests for ephemeral fallback when server rejects auth (401/403)."""

    def _make_batch_manager(self):
        """Create a TraceBatchManager with a pre-set trace_batch_id."""
        with patch(
            "crewai.events.listeners.tracing.trace_batch_manager.get_auth_token",
            return_value="mock_token",
        ):
            bm = TraceBatchManager()
        bm.current_batch = TraceBatch(
            user_context={"privacy_level": "standard"},
            execution_metadata={"execution_type": "crew", "crew_name": "test"},
        )
        bm.trace_batch_id = bm.current_batch.batch_id
        bm.is_current_batch_ephemeral = False
        return bm

    def test_401_non_ephemeral_falls_back_to_ephemeral(self):
        """A 401 on the non-ephemeral endpoint should retry as ephemeral."""
        bm = self._make_batch_manager()
        server_id = "ephemeral-fallback-id"

        auth_rejected = MagicMock(status_code=401, text="Bad credentials")
        ephemeral_success = MagicMock(
            status_code=201,
            json=MagicMock(return_value={"ephemeral_trace_id": server_id}),
        )

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_trace_batch",
                return_value=auth_rejected,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                return_value=ephemeral_success,
            ) as mock_ephemeral,
            patch("crewai.events.listeners.tracing.trace_batch_manager.time.sleep"),
        ):
            bm._initialize_backend_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={"execution_type": "crew"},
                use_ephemeral=False,
            )

        assert bm.trace_batch_id == server_id
        assert bm.is_current_batch_ephemeral is True
        mock_ephemeral.assert_called_once()

    def test_403_non_ephemeral_falls_back_to_ephemeral(self):
        """A 403 on the non-ephemeral endpoint should also fall back."""
        bm = self._make_batch_manager()
        server_id = "ephemeral-fallback-403"

        forbidden = MagicMock(status_code=403, text="Forbidden")
        ephemeral_success = MagicMock(
            status_code=201,
            json=MagicMock(return_value={"ephemeral_trace_id": server_id}),
        )

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_trace_batch",
                return_value=forbidden,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                return_value=ephemeral_success,
            ),
            patch("crewai.events.listeners.tracing.trace_batch_manager.time.sleep"),
        ):
            bm._initialize_backend_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={"execution_type": "crew"},
                use_ephemeral=False,
            )

        assert bm.trace_batch_id == server_id
        assert bm.is_current_batch_ephemeral is True

    def test_401_on_ephemeral_does_not_recurse(self):
        """A 401 on the ephemeral endpoint should NOT try to fall back again."""
        bm = self._make_batch_manager()
        bm.is_current_batch_ephemeral = True

        auth_rejected = MagicMock(status_code=401, text="Bad credentials")

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                return_value=auth_rejected,
            ) as mock_ephemeral,
            patch("crewai.events.listeners.tracing.trace_batch_manager.time.sleep"),
        ):
            bm._initialize_backend_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={"execution_type": "crew"},
                use_ephemeral=True,
            )

        assert bm.trace_batch_id is None
        # Called only once — no recursive fallback
        mock_ephemeral.assert_called()

    def test_401_fallback_ephemeral_also_fails(self):
        """If ephemeral fallback also fails, trace_batch_id is cleared."""
        bm = self._make_batch_manager()

        auth_rejected = MagicMock(status_code=401, text="Bad credentials")
        ephemeral_fail = MagicMock(status_code=422, text="Validation failed")

        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch.object(
                bm.plus_api,
                "initialize_trace_batch",
                return_value=auth_rejected,
            ),
            patch.object(
                bm.plus_api,
                "initialize_ephemeral_trace_batch",
                return_value=ephemeral_fail,
            ),
            patch("crewai.events.listeners.tracing.trace_batch_manager.time.sleep"),
        ):
            bm._initialize_backend_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={"execution_type": "crew"},
                use_ephemeral=False,
            )

        assert bm.trace_batch_id is None


class TestMarkBatchAsFailedRouting:
    """Tests: _mark_batch_as_failed routes to the correct endpoint."""

    def _make_batch_manager(self, ephemeral: bool = False):
        with patch(
            "crewai.events.listeners.tracing.trace_batch_manager.get_auth_token",
            return_value="mock_token",
        ):
            bm = TraceBatchManager()
        bm.is_current_batch_ephemeral = ephemeral
        return bm

    def test_routes_to_ephemeral_endpoint_when_ephemeral(self):
        """Ephemeral batches must use mark_ephemeral_trace_batch_as_failed."""
        bm = self._make_batch_manager(ephemeral=True)

        with patch.object(
            bm.plus_api, "mark_ephemeral_trace_batch_as_failed"
        ) as mock_ephemeral, patch.object(
            bm.plus_api, "mark_trace_batch_as_failed"
        ) as mock_non_ephemeral:
            bm._mark_batch_as_failed("batch-123", "some error")

        mock_ephemeral.assert_called_once_with("batch-123", "some error")
        mock_non_ephemeral.assert_not_called()

    def test_routes_to_non_ephemeral_endpoint_when_not_ephemeral(self):
        """Non-ephemeral batches must use mark_trace_batch_as_failed."""
        bm = self._make_batch_manager(ephemeral=False)

        with patch.object(
            bm.plus_api, "mark_ephemeral_trace_batch_as_failed"
        ) as mock_ephemeral, patch.object(
            bm.plus_api, "mark_trace_batch_as_failed"
        ) as mock_non_ephemeral:
            bm._mark_batch_as_failed("batch-456", "another error")

        mock_non_ephemeral.assert_called_once_with("batch-456", "another error")
        mock_ephemeral.assert_not_called()


class TestBackendInitializedGatedOnSuccess:
    """Tests: backend_initialized reflects actual init success on non-first-time path."""

    def test_backend_initialized_true_on_success(self):
        """backend_initialized is True when _initialize_backend_batch succeeds."""
        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.should_auto_collect_first_time_traces",
                return_value=False,
            ),
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.get_auth_token",
                return_value="mock_token",
            ),
        ):
            bm = TraceBatchManager()
            mock_response = MagicMock(
                status_code=201,
                json=MagicMock(return_value={"trace_id": "server-id"}),
            )
            with patch.object(
                bm.plus_api, "initialize_trace_batch", return_value=mock_response
            ):
                bm.initialize_batch(
                    user_context={"privacy_level": "standard"},
                    execution_metadata={"execution_type": "crew"},
                )

        assert bm.backend_initialized is True
        assert bm.trace_batch_id == "server-id"

    def test_backend_initialized_false_on_failure(self):
        """backend_initialized is False when _initialize_backend_batch fails."""
        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.should_auto_collect_first_time_traces",
                return_value=False,
            ),
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.get_auth_token",
                return_value="mock_token",
            ),
        ):
            bm = TraceBatchManager()
            with patch.object(
                bm.plus_api, "initialize_trace_batch", return_value=None
            ):
                bm.initialize_batch(
                    user_context={"privacy_level": "standard"},
                    execution_metadata={"execution_type": "crew"},
                )

        assert bm.backend_initialized is False
        assert bm.trace_batch_id is None


class TestTraceBatchManagerDuplicateInitMerge:
    """Second initialize_batch call merges execution_metadata (flow after lazy action)."""

    def test_duplicate_initialize_merges_execution_metadata(self):
        with (
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.should_auto_collect_first_time_traces",
                return_value=True,
            ),
            patch(
                "crewai.events.listeners.tracing.trace_batch_manager.is_tracing_enabled_in_context",
                return_value=True,
            ),
        ):
            bm = TraceBatchManager()
            bm.initialize_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={
                    "crew_name": "Unknown Crew",
                    "crewai_version": "9.9.9",
                },
            )
            first_batch_id = bm.current_batch.batch_id
            bm.initialize_batch(
                user_context={"privacy_level": "standard"},
                execution_metadata={
                    "flow_name": "ResearchFlow",
                    "execution_type": "flow",
                    "crewai_version": "9.9.9",
                    "execution_start": "2026-01-01T00:00:00+00:00",
                },
            )

        assert bm.current_batch.batch_id == first_batch_id
        meta = bm.current_batch.execution_metadata
        assert meta.get("execution_type") == "flow"
        assert meta.get("flow_name") == "ResearchFlow"
        assert meta.get("crew_name") == "Unknown Crew"