"""Tests for LangChain integration."""

from unittest.mock import MagicMock, patch
from uuid import uuid4

import pytest

from lumenova_beacon import BeaconClient, BeaconLangGraphHandler
from lumenova_beacon.types import SpanType, SpanKind, StatusCode


@pytest.fixture
def mock_client():
    """Create a mock BeaconClient."""
    client = MagicMock(spec=BeaconClient)
    client.should_sample = MagicMock(return_value=True)
    return client


@pytest.fixture
def handler(mock_client):
    """Create a BeaconLangGraphHandler with mock client."""
    with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
        mock_get.return_value = mock_client
        return BeaconLangGraphHandler()


class TestBeaconLangGraphHandler:
    """Test suite for BeaconLangGraphHandler."""

    def test_init_with_defaults(self):
        """Test initialization with default parameters."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            mock_client = MagicMock(spec=BeaconClient)
            mock_get.return_value = mock_client
            handler = BeaconLangGraphHandler()
            assert handler.client == mock_client
            assert handler._runs == {}
            assert handler._session_id is None

    def test_init_with_custom_parameters(self):
        """Test initialization with custom session_id."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            mock_client = MagicMock(spec=BeaconClient)
            mock_get.return_value = mock_client
            handler = BeaconLangGraphHandler(
                session_id="test-session"
            )
            assert handler.client is not None
            assert handler._session_id == "test-session"
            mock_get.assert_called_once()

    def test_on_chain_start(self, handler, mock_client):
        """Test chain start event creates span."""
        run_id = uuid4()
        serialized = {"name": "TestChain", "id": ["test", "chain"]}
        inputs = {"query": "test question"}
        metadata = {"user_id": "123"}
        tags = ["test", "chain"]

        handler.on_chain_start(
            serialized=serialized,
            inputs=inputs,
            run_id=run_id,
            metadata=metadata,
            tags=tags,
        )

        # Check run was created
        run_id_str = str(run_id)
        assert run_id_str in handler._runs
        run_data = handler._runs[run_id_str]

        # Check span properties
        span = run_data["span"]
        assert span.name == "TestChain"
        assert span.span_type == SpanType.CHAIN
        assert span.kind == SpanKind.INTERNAL
        assert span.start_time is not None

        # Check attributes
        assert "gen_ai.input.messages" in span.attributes
        assert "langchain.metadata.user_id" in span.attributes
        assert "langchain.metadata.tags" in span.attributes

    def test_on_chain_end(self, handler, mock_client):
        """Test chain end event finalizes and sends span."""
        run_id = uuid4()

        # Start chain first
        handler.on_chain_start(
            serialized={"name": "TestChain"},
            inputs={"query": "test"},
            run_id=run_id,
        )

        # End chain
        outputs = {"result": "test answer"}
        handler.on_chain_end(outputs=outputs, run_id=run_id)

        # Check span was sent
        mock_client.export_span.assert_called_once()
        sent_span = mock_client.export_span.call_args[0][0]

        # Check span properties
        assert sent_span.end_time is not None
        assert sent_span.status_code == StatusCode.OK
        assert "gen_ai.output.messages" in sent_span.attributes

        # Check run was cleaned up
        assert str(run_id) not in handler._runs

    def test_on_chain_error(self, handler, mock_client):
        """Test chain error event records exception."""
        run_id = uuid4()

        # Start chain first
        handler.on_chain_start(
            serialized={"name": "TestChain"},
            inputs={"query": "test"},
            run_id=run_id,
        )

        # Error
        error = ValueError("Test error")
        handler.on_chain_error(error=error, run_id=run_id)

        # Check span was sent with error
        mock_client.export_span.assert_called_once()
        sent_span = mock_client.export_span.call_args[0][0]

        # Check error was recorded
        assert sent_span.status_code == StatusCode.ERROR
        assert sent_span.status_description == "Test error"
        assert "exception.type" in sent_span.attributes
        assert "exception.message" in sent_span.attributes

    def test_on_llm_start(self, handler, mock_client):
        """Test LLM start event creates generation span."""
        run_id = uuid4()
        serialized = {
            "name": "ChatOpenAI",
            "kwargs": {"model_name": "gpt-4"},
        }
        prompts = ["What is the capital of France?"]

        handler.on_llm_start(
            serialized=serialized,
            prompts=prompts,
            run_id=run_id,
        )

        # Check run was created
        run_id_str = str(run_id)
        assert run_id_str in handler._runs
        span = handler._runs[run_id_str]["span"]

        # Check span properties
        assert span.span_type == SpanType.GENERATION
        assert span.kind == SpanKind.CLIENT
        assert "gen_ai.input.messages" in span.attributes
        assert "gen_ai.request.model" in span.attributes
        assert span.attributes["gen_ai.request.model"] == "gpt-4"

    def test_on_chat_model_start(self, handler, mock_client):
        """Test chat model start event."""
        from unittest.mock import MagicMock

        run_id = uuid4()
        serialized = {"name": "ChatOpenAI"}

        # Mock message objects
        message1 = MagicMock()
        message1.type = "human"
        message1.content = "Hello"

        messages = [[message1]]

        handler.on_chat_model_start(
            serialized=serialized,
            messages=messages,
            run_id=run_id,
        )

        # Check span was created
        assert str(run_id) in handler._runs
        span = handler._runs[str(run_id)]["span"]
        assert span.span_type == SpanType.GENERATION

    def test_on_llm_end_with_token_usage(self, handler, mock_client):
        """Test LLM end event extracts token usage."""
        from unittest.mock import MagicMock

        run_id = uuid4()

        # Start LLM first
        handler.on_llm_start(
            serialized={"name": "ChatOpenAI"},
            prompts=["test"],
            run_id=run_id,
        )

        # Mock LLMResult with token usage
        response = MagicMock()
        generation = MagicMock()
        generation.text = "Paris"
        generation.message = None
        response.generations = [[generation]]
        response.llm_output = {
            "token_usage": {
                "prompt_tokens": 10,
                "completion_tokens": 5,
                "total_tokens": 15,
            },
            "model_name": "gpt-4",
        }

        handler.on_llm_end(response=response, run_id=run_id)

        # Check token usage was captured as individual attributes
        sent_span = mock_client.export_span.call_args[0][0]
        assert "gen_ai.usage.input_tokens" in sent_span.attributes
        assert "gen_ai.usage.output_tokens" in sent_span.attributes

        # Verify token counts
        assert sent_span.attributes["gen_ai.usage.input_tokens"] == 10
        assert sent_span.attributes["gen_ai.usage.output_tokens"] == 5

        # Verify model name was captured
        assert "gen_ai.request.model" in sent_span.attributes
        assert sent_span.attributes["gen_ai.request.model"] == "gpt-4"

    def test_on_tool_start(self, handler, mock_client):
        """Test tool start event."""
        run_id = uuid4()
        serialized = {"name": "Calculator"}
        input_str = "2 + 2"

        handler.on_tool_start(
            serialized=serialized,
            input_str=input_str,
            run_id=run_id,
        )

        # Check span was created
        assert str(run_id) in handler._runs
        span = handler._runs[str(run_id)]["span"]
        assert span.span_type == SpanType.TOOL
        assert span.kind == SpanKind.INTERNAL

    def test_on_tool_end(self, handler, mock_client):
        """Test tool end event."""
        run_id = uuid4()

        # Start tool first
        handler.on_tool_start(
            serialized={"name": "Calculator"},
            input_str="2 + 2",
            run_id=run_id,
        )

        # End tool
        output = "4"
        handler.on_tool_end(output=output, run_id=run_id)

        # Check span was sent
        mock_client.export_span.assert_called_once()
        sent_span = mock_client.export_span.call_args[0][0]
        assert "gen_ai.output.messages" in sent_span.attributes

    def test_on_retriever_start(self, handler, mock_client):
        """Test retriever start event."""
        run_id = uuid4()
        serialized = {"name": "VectorStoreRetriever"}
        query = "test query"

        handler.on_retriever_start(
            serialized=serialized,
            query=query,
            run_id=run_id,
        )

        # Check span was created
        assert str(run_id) in handler._runs
        span = handler._runs[str(run_id)]["span"]
        assert span.span_type == SpanType.RETRIEVAL

    def test_on_retriever_end(self, handler, mock_client):
        """Test retriever end event."""
        from unittest.mock import MagicMock

        run_id = uuid4()

        # Start retriever first
        handler.on_retriever_start(
            serialized={"name": "VectorStoreRetriever"},
            query="test",
            run_id=run_id,
        )

        # Mock documents
        doc1 = MagicMock()
        doc1.page_content = "Document 1"
        doc1.metadata = {"source": "test.txt"}

        documents = [doc1]

        handler.on_retriever_end(documents=documents, run_id=run_id)

        # Check span was sent
        sent_span = mock_client.export_span.call_args[0][0]
        assert "gen_ai.output.messages" in sent_span.attributes
        assert "retriever.document_count" in sent_span.attributes

    def test_parent_child_relationship(self, handler, mock_client):
        """Test parent-child span relationships."""
        parent_run_id = uuid4()
        child_run_id = uuid4()

        # Start parent chain
        handler.on_chain_start(
            serialized={"name": "ParentChain"},
            inputs={"query": "test"},
            run_id=parent_run_id,
        )

        # Start child LLM
        handler.on_llm_start(
            serialized={"name": "ChildLLM"},
            prompts=["test"],
            run_id=child_run_id,
            parent_run_id=parent_run_id,
        )

        # Check parent-child relationship
        parent_span = handler._runs[str(parent_run_id)]["span"]
        child_span = handler._runs[str(child_run_id)]["span"]

        assert child_span.parent_id == parent_span.span_id
        assert child_span.trace_id == parent_span.trace_id

    def test_on_agent_action(self, handler, mock_client):
        """Test agent action event."""
        from unittest.mock import MagicMock

        parent_run_id = uuid4()
        action_run_id = uuid4()

        # Start parent agent chain
        handler.on_chain_start(
            serialized={"name": "AgentExecutor"},
            inputs={"input": "test"},
            run_id=parent_run_id,
        )

        # Record agent action
        action = MagicMock()
        action.tool = "Calculator"
        action.tool_input = "2 + 2"
        action.log = "Thought: I need to calculate 2 + 2"

        handler.on_agent_action(
            action=action,
            run_id=action_run_id,
            parent_run_id=parent_run_id,
        )

        # Check action was recorded in parent span
        parent_span = handler._runs[str(parent_run_id)]["span"]
        action_key = f"agent.action.{str(action_run_id)[:8]}"
        assert action_key in parent_span.attributes

    def test_on_agent_finish(self, handler, mock_client):
        """Test agent finish event."""
        from unittest.mock import MagicMock

        parent_run_id = uuid4()
        finish_run_id = uuid4()

        # Start parent agent chain
        handler.on_chain_start(
            serialized={"name": "AgentExecutor"},
            inputs={"input": "test"},
            run_id=parent_run_id,
        )

        # Record agent finish
        finish = MagicMock()
        finish.return_values = {"output": "4"}
        finish.log = "Final Answer: 4"

        handler.on_agent_finish(
            finish=finish,
            run_id=finish_run_id,
            parent_run_id=parent_run_id,
        )

        # Check finish was recorded in parent span
        parent_span = handler._runs[str(parent_run_id)]["span"]
        assert "agent.finish" in parent_span.attributes

    def test_defensive_error_handling(self, handler, mock_client):
        """Test that handler doesn't crash on errors."""
        # Try to end a run that doesn't exist - should not raise
        handler.on_chain_end(outputs={}, run_id=uuid4())

        # Try to error a run that doesn't exist - should not raise
        handler.on_chain_error(error=ValueError("test"), run_id=uuid4())

        # Handler should still be functional
        run_id = uuid4()
        handler.on_chain_start(
            serialized={"name": "TestChain"},
            inputs={},
            run_id=run_id,
        )
        assert str(run_id) in handler._runs

    def test_extract_component_type(self, handler):
        """Test component type extraction from serialized data."""
        # Test with name
        serialized = {"name": "ChatOpenAI"}
        assert handler._extract_component_type(serialized) == "ChatOpenAI"

        # Test with id list
        serialized = {"id": ["langchain", "chains", "RetrievalQA"]}
        assert handler._extract_component_type(serialized) == "RetrievalQA"

        # Test with unknown format
        serialized = {}
        assert handler._extract_component_type(serialized) == "Unknown"

    def test_nested_tool_and_function_spans(self, handler, mock_client):
        """Test that @trace decorated functions create nested spans under tool spans."""
        from lumenova_beacon import trace
        from lumenova_beacon.tracing.trace import get_current_span

        tool_run_id = uuid4()

        # Start tool span (simulating LangChain's on_tool_start)
        handler.on_tool_start(
            serialized={"name": "TestTool"},
            input_str="test input",
            run_id=tool_run_id,
        )

        # Verify tool span is set as current
        tool_span = handler._runs[str(tool_run_id)]["span"]
        current_span = get_current_span()
        assert current_span == tool_span
        assert tool_span.span_type == SpanType.TOOL

        # Simulate @trace decorated function executing inside the tool
        @trace()
        def tool_implementation(x: int) -> int:
            """Simulated tool function with @trace decorator."""
            return x * 2

        result = tool_implementation(5)
        assert result == 10

        # The @trace decorator should have created a function span
        # and it should have been sent when the function completed
        assert mock_client.export_span.call_count >= 1

        # Get the function span (should be the first call)
        function_span = mock_client.export_span.call_args_list[0][0][0]

        # Verify function span properties
        assert function_span.span_type == SpanType.FUNCTION
        assert function_span.name == "tool_implementation"
        assert function_span.parent_id == tool_span.span_id
        assert function_span.trace_id == tool_span.trace_id

        # End tool span
        handler.on_tool_end(output="test output", run_id=tool_run_id)

        # Verify tool span was sent
        assert mock_client.export_span.call_count >= 2
        sent_tool_span = mock_client.export_span.call_args_list[-1][0][0]
        assert sent_tool_span.span_type == SpanType.TOOL

    def test_context_restoration_after_tool_completion(self, handler, mock_client):
        """Test that span context is properly restored after tool completion."""
        from lumenova_beacon.tracing.trace import get_current_span

        parent_run_id = uuid4()
        tool_run_id = uuid4()

        # Start parent chain
        handler.on_chain_start(
            serialized={"name": "ParentChain"},
            inputs={"query": "test"},
            run_id=parent_run_id,
        )

        parent_span = handler._runs[str(parent_run_id)]["span"]
        assert get_current_span() == parent_span

        # Start tool (child of chain)
        handler.on_tool_start(
            serialized={"name": "TestTool"},
            input_str="test",
            run_id=tool_run_id,
            parent_run_id=parent_run_id,
        )

        tool_span = handler._runs[str(tool_run_id)]["span"]
        assert get_current_span() == tool_span
        assert tool_span.parent_id == parent_span.span_id

        # End tool
        handler.on_tool_end(output="result", run_id=tool_run_id)

        # Context should be restored to parent
        assert get_current_span() == parent_span

        # End parent
        handler.on_chain_end(outputs={"result": "done"}, run_id=parent_run_id)

    def test_context_restoration_on_error(self, handler, mock_client):
        """Test that span context is properly restored even when tool errors."""
        from lumenova_beacon.tracing.trace import get_current_span

        parent_run_id = uuid4()
        tool_run_id = uuid4()

        # Start parent chain
        handler.on_chain_start(
            serialized={"name": "ParentChain"},
            inputs={"query": "test"},
            run_id=parent_run_id,
        )

        parent_span = handler._runs[str(parent_run_id)]["span"]

        # Start tool
        handler.on_tool_start(
            serialized={"name": "TestTool"},
            input_str="test",
            run_id=tool_run_id,
            parent_run_id=parent_run_id,
        )

        # Tool errors
        error = ValueError("Tool failed")
        handler.on_tool_error(error=error, run_id=tool_run_id, parent_run_id=parent_run_id)

        # Context should be restored to parent even after error
        assert get_current_span() == parent_span

        # End parent
        handler.on_chain_end(outputs={"result": "done"}, run_id=parent_run_id)

    def test_multiple_nested_function_calls_in_tool(self, handler, mock_client):
        """Test multiple @trace decorated functions within a single tool."""
        from lumenova_beacon import trace

        tool_run_id = uuid4()

        # Start tool span
        handler.on_tool_start(
            serialized={"name": "ComplexTool"},
            input_str="test",
            run_id=tool_run_id,
        )

        tool_span = handler._runs[str(tool_run_id)]["span"]

        # Define multiple traced helper functions
        @trace()
        def helper1(x: int) -> int:
            return x + 1

        @trace()
        def helper2(x: int) -> int:
            return x * 2

        # Execute them
        result1 = helper1(5)
        result2 = helper2(result1)

        assert result1 == 6
        assert result2 == 12

        # Both function spans should have been sent
        # Each should have the tool span as parent and share trace_id
        function_calls = [call for call in mock_client.export_span.call_args_list]
        assert len(function_calls) >= 2

        # Check both function spans
        helper1_span = function_calls[0][0][0]
        helper2_span = function_calls[1][0][0]

        assert helper1_span.span_type == SpanType.FUNCTION
        assert helper1_span.parent_id == tool_span.span_id
        assert helper1_span.trace_id == tool_span.trace_id

        assert helper2_span.span_type == SpanType.FUNCTION
        assert helper2_span.parent_id == tool_span.span_id
        assert helper2_span.trace_id == tool_span.trace_id

        # End tool
        handler.on_tool_end(output="done", run_id=tool_run_id)

    def test_context_cleared_after_root_langchain_span(self, handler, mock_client):
        """Test that context is cleared after root LangChain span completes.

        This prevents subsequent @trace calls from incorrectly nesting under
        the completed LangChain span.
        """
        from lumenova_beacon import trace
        from lumenova_beacon.tracing.trace import get_current_span

        # Execute a root LangChain chain (no parent)
        chain_run_id = uuid4()
        handler.on_chain_start(
            serialized={"name": "RootChain"},
            inputs={"query": "test"},
            run_id=chain_run_id,
        )

        chain_span = handler._runs[str(chain_run_id)]["span"]
        assert get_current_span() == chain_span

        # End the chain
        handler.on_chain_end(outputs={"result": "done"}, run_id=chain_run_id)

        # Context should be cleared (no current span)
        assert get_current_span() is None

        # Now trace a completely separate function
        @trace()
        def separate_function() -> str:
            return "independent"

        result = separate_function()
        assert result == "independent"

        # Get the function span (should be sent to mock_client)
        function_span = mock_client.export_span.call_args_list[-1][0][0]

        # Verify the function span is NOT a child of the chain span
        assert function_span.span_type == SpanType.FUNCTION
        assert function_span.name == "separate_function"
        assert function_span.parent_id is None  # No parent!
        assert function_span.trace_id != chain_span.trace_id  # Different trace!


class TestBeaconLangGraphHandlerMetadata:
    """Test suite for BeaconLangGraphHandler metadata functionality."""

    def test_init_with_environment(self):
        """Test initialization with environment parameter."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            mock_client = MagicMock(spec=BeaconClient)
            mock_client.config.session_id = None
            mock_get.return_value = mock_client
            handler = BeaconLangGraphHandler(environment="production")
            assert handler._environment == "production"

    def test_init_with_agent_name(self):
        """Test initialization with agent_name parameter."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            mock_client = MagicMock(spec=BeaconClient)
            mock_client.config.session_id = None
            mock_get.return_value = mock_client
            handler = BeaconLangGraphHandler(agent_name="my-agent")
            assert handler._agent_name == "my-agent"

    def test_init_with_metadata(self):
        """Test initialization with custom metadata dict."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            mock_client = MagicMock(spec=BeaconClient)
            mock_client.config.session_id = None
            mock_get.return_value = mock_client
            handler = BeaconLangGraphHandler(
                metadata={"app_name": "my-app", "version": "1.0.0"}
            )
            assert handler._metadata == {"app_name": "my-app", "version": "1.0.0"}

    def test_init_with_all_new_parameters(self):
        """Test initialization with all new metadata parameters."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            mock_client = MagicMock(spec=BeaconClient)
            mock_client.config.session_id = None
            mock_get.return_value = mock_client
            handler = BeaconLangGraphHandler(
                session_id="test-session",
                environment="staging",
                agent_name="test-agent",
                metadata={"team": "platform", "priority": 1}
            )
            assert handler._session_id == "test-session"
            assert handler._environment == "staging"
            assert handler._agent_name == "test-agent"
            assert handler._metadata == {"team": "platform", "priority": 1}

    def test_metadata_defaults_to_empty_dict(self):
        """Test that metadata defaults to empty dict when not provided."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            mock_client = MagicMock(spec=BeaconClient)
            mock_client.config.session_id = None
            mock_get.return_value = mock_client
            handler = BeaconLangGraphHandler()
            assert handler._metadata == {}

    def test_environment_attribute_set_on_span(self):
        """Test that environment is set as span attribute."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            with patch("lumenova_beacon.tracing.integrations.langchain.trace") as mock_trace:
                mock_client = MagicMock(spec=BeaconClient)
                mock_client.config.session_id = None
                mock_get.return_value = mock_client

                mock_span = MagicMock()
                mock_tracer = MagicMock()
                mock_tracer.start_span.return_value = mock_span
                mock_trace.get_tracer.return_value = mock_tracer
                mock_trace.set_span_in_context.return_value = MagicMock()

                handler = BeaconLangGraphHandler(environment="production")

                run_id = uuid4()
                handler.on_chain_start(
                    serialized={"name": "TestChain"},
                    inputs={"query": "test"},
                    run_id=run_id,
                )

                # Verify environment attribute was set (OTEL standard)
                mock_span.set_attribute.assert_any_call("deployment.environment.name", "production")

    def test_agent_name_attribute_set_on_span(self):
        """Test that agent_name is set as span attribute."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            with patch("lumenova_beacon.tracing.integrations.langchain.trace") as mock_trace:
                mock_client = MagicMock(spec=BeaconClient)
                mock_client.config.session_id = None
                mock_get.return_value = mock_client

                mock_span = MagicMock()
                mock_tracer = MagicMock()
                mock_tracer.start_span.return_value = mock_span
                mock_trace.get_tracer.return_value = mock_tracer
                mock_trace.set_span_in_context.return_value = MagicMock()

                handler = BeaconLangGraphHandler(agent_name="my-agent")

                run_id = uuid4()
                handler.on_chain_start(
                    serialized={"name": "TestChain"},
                    inputs={"query": "test"},
                    run_id=run_id,
                )

                # Verify agent_name attribute was set (OTEL standard)
                mock_span.set_attribute.assert_any_call("gen_ai.agent.name", "my-agent")

    def test_custom_metadata_attributes_set_on_span(self):
        """Test that custom metadata dict values are set as span attributes."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            with patch("lumenova_beacon.tracing.integrations.langchain.trace") as mock_trace:
                mock_client = MagicMock(spec=BeaconClient)
                mock_client.config.session_id = None
                mock_get.return_value = mock_client

                mock_span = MagicMock()
                mock_tracer = MagicMock()
                mock_tracer.start_span.return_value = mock_span
                mock_trace.get_tracer.return_value = mock_tracer
                mock_trace.set_span_in_context.return_value = MagicMock()

                handler = BeaconLangGraphHandler(
                    metadata={"app_name": "my-app", "version": "1.0.0", "count": 42}
                )

                run_id = uuid4()
                handler.on_chain_start(
                    serialized={"name": "TestChain"},
                    inputs={"query": "test"},
                    run_id=run_id,
                )

                # Verify custom metadata attributes were set with beacon.metadata. prefix
                mock_span.set_attribute.assert_any_call("beacon.metadata.app_name", "my-app")
                mock_span.set_attribute.assert_any_call("beacon.metadata.version", "1.0.0")
                mock_span.set_attribute.assert_any_call("beacon.metadata.count", 42)

    def test_complex_metadata_value_json_serialized(self):
        """Test that complex metadata values are JSON serialized."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            with patch("lumenova_beacon.tracing.integrations.langchain.trace") as mock_trace:
                mock_client = MagicMock(spec=BeaconClient)
                mock_client.config.session_id = None
                mock_get.return_value = mock_client

                mock_span = MagicMock()
                mock_tracer = MagicMock()
                mock_tracer.start_span.return_value = mock_span
                mock_trace.get_tracer.return_value = mock_tracer
                mock_trace.set_span_in_context.return_value = MagicMock()

                handler = BeaconLangGraphHandler(
                    metadata={"tags": ["tag1", "tag2"], "config": {"key": "value"}}
                )

                run_id = uuid4()
                handler.on_chain_start(
                    serialized={"name": "TestChain"},
                    inputs={"query": "test"},
                    run_id=run_id,
                )

                # Verify complex values are JSON serialized
                import json
                mock_span.set_attribute.assert_any_call(
                    "beacon.metadata.tags", json.dumps(["tag1", "tag2"])
                )
                mock_span.set_attribute.assert_any_call(
                    "beacon.metadata.config", json.dumps({"key": "value"})
                )

    def test_langgraph_graph_name_extraction(self):
        """Test that graph name is extracted from LangGraph metadata when agent_name not set."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            with patch("lumenova_beacon.tracing.integrations.langchain.trace") as mock_trace:
                mock_client = MagicMock(spec=BeaconClient)
                mock_client.config.session_id = None
                mock_get.return_value = mock_client

                mock_span = MagicMock()
                mock_tracer = MagicMock()
                mock_tracer.start_span.return_value = mock_span
                mock_trace.get_tracer.return_value = mock_tracer
                mock_trace.set_span_in_context.return_value = MagicMock()

                # No agent_name set - should extract from LangGraph metadata
                handler = BeaconLangGraphHandler()

                run_id = uuid4()
                handler.on_chain_start(
                    serialized={"name": "TestChain"},
                    inputs={"query": "test"},
                    run_id=run_id,
                    metadata={"langgraph_checkpoint_ns": "my_graph_name"},
                )

                # Verify graph name was extracted and set (OTEL standard)
                mock_span.set_attribute.assert_any_call("gen_ai.agent.name", "my_graph_name")

    def test_agent_name_not_overridden_by_langgraph(self):
        """Test that explicit agent_name is not overridden by LangGraph metadata."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            with patch("lumenova_beacon.tracing.integrations.langchain.trace") as mock_trace:
                mock_client = MagicMock(spec=BeaconClient)
                mock_client.config.session_id = None
                mock_get.return_value = mock_client

                mock_span = MagicMock()
                mock_tracer = MagicMock()
                mock_tracer.start_span.return_value = mock_span
                mock_trace.get_tracer.return_value = mock_tracer
                mock_trace.set_span_in_context.return_value = MagicMock()

                # Explicit agent_name set - should NOT be overridden
                handler = BeaconLangGraphHandler(agent_name="my-explicit-agent")

                run_id = uuid4()
                handler.on_chain_start(
                    serialized={"name": "TestChain"},
                    inputs={"query": "test"},
                    run_id=run_id,
                    metadata={"langgraph_checkpoint_ns": "graph_name_from_langgraph"},
                )

                # Verify explicit agent_name was used (set in _start_span) (OTEL standard)
                mock_span.set_attribute.assert_any_call("gen_ai.agent.name", "my-explicit-agent")

                # Verify graph name was NOT set (no second call with graph name)
                calls = [call for call in mock_span.set_attribute.call_args_list
                         if call[0][0] == "gen_ai.agent.name"]
                assert len(calls) == 1
                assert calls[0][0][1] == "my-explicit-agent"

    def test_agent_id_attribute_set_on_span(self):
        """Test that agent_id is set as span attribute."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            with patch("lumenova_beacon.tracing.integrations.langchain.trace") as mock_trace:
                mock_client = MagicMock(spec=BeaconClient)
                mock_client.config.session_id = None
                mock_get.return_value = mock_client

                mock_span = MagicMock()
                mock_tracer = MagicMock()
                mock_tracer.start_span.return_value = mock_span
                mock_trace.get_tracer.return_value = mock_tracer
                mock_trace.set_span_in_context.return_value = MagicMock()

                handler = BeaconLangGraphHandler(agent_id="agent-001")

                run_id = uuid4()
                handler.on_chain_start(
                    serialized={"name": "TestChain"},
                    inputs={"query": "test"},
                    run_id=run_id,
                )

                # Verify agent_id attribute was set (OTEL standard)
                mock_span.set_attribute.assert_any_call("gen_ai.agent.id", "agent-001")

    def test_agent_description_attribute_set_on_span(self):
        """Test that agent_description is set as span attribute."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            with patch("lumenova_beacon.tracing.integrations.langchain.trace") as mock_trace:
                mock_client = MagicMock(spec=BeaconClient)
                mock_client.config.session_id = None
                mock_get.return_value = mock_client

                mock_span = MagicMock()
                mock_tracer = MagicMock()
                mock_tracer.start_span.return_value = mock_span
                mock_trace.get_tracer.return_value = mock_tracer
                mock_trace.set_span_in_context.return_value = MagicMock()

                handler = BeaconLangGraphHandler(agent_description="A helpful assistant")

                run_id = uuid4()
                handler.on_chain_start(
                    serialized={"name": "TestChain"},
                    inputs={"query": "test"},
                    run_id=run_id,
                )

                # Verify agent_description attribute was set (OTEL standard)
                mock_span.set_attribute.assert_any_call("gen_ai.agent.description", "A helpful assistant")

    def test_model_parameters_extracted_on_llm_start(self):
        """Test that model parameters are extracted and set on LLM spans."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            with patch("lumenova_beacon.tracing.integrations.langchain.trace") as mock_trace:
                mock_client = MagicMock(spec=BeaconClient)
                mock_client.config.session_id = None
                mock_get.return_value = mock_client

                mock_span = MagicMock()
                mock_tracer = MagicMock()
                mock_tracer.start_span.return_value = mock_span
                mock_trace.get_tracer.return_value = mock_tracer
                mock_trace.set_span_in_context.return_value = MagicMock()

                handler = BeaconLangGraphHandler()

                run_id = uuid4()
                serialized = {
                    "name": "ChatOpenAI",
                    "kwargs": {
                        "model_name": "gpt-4",
                        "temperature": 0.7,
                        "top_p": 0.9,
                        "max_tokens": 1000,
                    }
                }

                handler.on_llm_start(
                    serialized=serialized,
                    prompts=["Hello"],
                    run_id=run_id,
                )

                # Verify model parameters were set (OTEL standard)
                mock_span.set_attribute.assert_any_call("gen_ai.request.temperature", 0.7)
                mock_span.set_attribute.assert_any_call("gen_ai.request.top_p", 0.9)
                mock_span.set_attribute.assert_any_call("gen_ai.request.max_tokens", 1000)


class TestMarkInterrupt:
    """Tests for mark_interrupt() — breakpoint interrupt detection."""

    @pytest.fixture
    def handler_with_agent(self):
        """Handler with agent_name and mock client."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            mock_client = MagicMock(spec=BeaconClient)
            mock_client.config.session_id = None
            mock_get.return_value = mock_client
            handler = BeaconLangGraphHandler(agent_name="task_planner")
            yield handler, mock_client

    def test_mark_interrupt_before_root_ends(self, handler_with_agent):
        """mark_interrupt() called while root span is still active propagates via on_chain_end."""
        handler, mock_client = handler_with_agent
        root_run_id = uuid4()

        # Start root chain
        handler.on_chain_start(
            serialized={"name": "LangGraph"},
            inputs={"messages": []},
            run_id=root_run_id,
        )

        # Simulate breakpoint detected in astream_events before root ends
        handler.mark_interrupt(value="Review research findings", interrupt_id="ckpt-123")

        assert handler.interrupt_detected is True
        assert handler._interrupt_info["value"] == "Review research findings"
        assert handler._interrupt_info["id"] == "ckpt-123"

        # Root span ends normally (on_chain_end, not on_chain_error)
        handler.on_chain_end(outputs={"result": "done"}, run_id=root_run_id)

        # The root span should have interrupt attributes (set by on_chain_end propagation)
        sent_span = mock_client.export_span.call_args[0][0]
        assert sent_span.attributes.get("langgraph.interrupt") is True
        assert sent_span.attributes.get("langgraph.interrupt.value") == "Review research findings"
        assert sent_span.attributes.get("langgraph.interrupt.id") == "ckpt-123"
        assert sent_span.status_code == StatusCode.OK

    def test_mark_interrupt_after_root_ends(self, handler_with_agent):
        """mark_interrupt() called after root span already exported triggers re-export."""
        handler, mock_client = handler_with_agent
        root_run_id = uuid4()

        # Start and end root chain normally
        handler.on_chain_start(
            serialized={"name": "LangGraph"},
            inputs={"messages": []},
            run_id=root_run_id,
        )
        handler.on_chain_end(outputs={"result": "done"}, run_id=root_run_id)

        # First export: normal completion (no interrupt)
        assert mock_client.export_span.call_count == 1
        first_span = mock_client.export_span.call_args_list[0][0][0]
        assert first_span.attributes.get("langgraph.interrupt") is not True

        # Now mark_interrupt called after root ended (late detection)
        handler.mark_interrupt(value="Awaiting approval", interrupt_id="ckpt-456")

        # Should trigger a re-export with interrupt attributes
        assert mock_client.export_span.call_count == 2
        reexported_span = mock_client.export_span.call_args_list[1][0][0]
        assert reexported_span.span_id == handler.root_span_id
        assert reexported_span.trace_id == handler.trace_id
        assert reexported_span.attributes.get("langgraph.interrupt") is True
        assert reexported_span.attributes.get("langgraph.interrupt.value") == "Awaiting approval"
        assert reexported_span.attributes.get("langgraph.interrupt.id") == "ckpt-456"
        assert reexported_span.attributes.get("gen_ai.agent.name") == "task_planner"

    def test_mark_interrupt_defaults_value(self, handler_with_agent):
        """mark_interrupt() without value defaults to 'Awaiting user input'."""
        handler, mock_client = handler_with_agent
        root_run_id = uuid4()

        handler.on_chain_start(
            serialized={"name": "LangGraph"},
            inputs={"messages": []},
            run_id=root_run_id,
        )
        handler.on_chain_end(outputs={}, run_id=root_run_id)

        handler.mark_interrupt()

        assert handler._interrupt_info["value"] == "Awaiting user input"
        assert handler._interrupt_info["id"] is None

    def test_mark_interrupt_persists_id_for_resume(self, handler_with_agent):
        """mark_interrupt() persists interrupt_id at class level for resume handlers."""
        handler, _ = handler_with_agent
        root_run_id = uuid4()

        handler.on_chain_start(
            serialized={"name": "LangGraph"},
            inputs={"messages": []},
            run_id=root_run_id,
        )

        handler.mark_interrupt(value="Review", interrupt_id="ckpt-789")

        # Verify class-level storage
        trace_id = handler.trace_id
        assert BeaconLangGraphHandler._trace_interrupt_ids.get(trace_id) == "ckpt-789"

        # Verify interrupt_id property returns it
        assert handler.interrupt_id == "ckpt-789"

        # Clean up class-level state
        BeaconLangGraphHandler._trace_interrupt_ids.pop(trace_id, None)

    def test_interrupt_detected_property(self, handler_with_agent):
        """interrupt_detected property reflects mark_interrupt state."""
        handler, _ = handler_with_agent

        assert handler.interrupt_detected is False
        handler.mark_interrupt(value="test")
        assert handler.interrupt_detected is True


class TestConditionalEdgeDetection:
    """Test detection of LangGraph conditional routing functions."""

    @pytest.fixture
    def handler_with_mock(self):
        """Create a handler with mock client for conditional edge tests."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            mock_client = MagicMock(spec=BeaconClient)
            mock_client.should_sample = MagicMock(return_value=True)
            mock_get.return_value = mock_client
            handler = BeaconLangGraphHandler()
            yield handler, mock_client

    def test_conditional_edge_with_string_output(self, handler_with_mock):
        """Routing function with seq:step tag and string output is classified as CONDITIONAL."""
        handler, mock_client = handler_with_mock
        run_id = uuid4()

        handler.on_chain_start(
            serialized=None,
            inputs={"messages": [{"role": "user", "content": "hello"}]},
            run_id=run_id,
            metadata={"langgraph_node": "agent", "langgraph_checkpoint_ns": "graph"},
            tags=["seq:step:3"],
            name="should_continue",
        )

        assert str(run_id) in handler._potential_conditional_run_ids

        handler.on_chain_end(
            outputs="tools",
            run_id=run_id,
        )

        assert str(run_id) not in handler._potential_conditional_run_ids

    def test_conditional_edge_captures_routing_decision(self, handler_with_mock):
        """Routing decision is captured as a span attribute."""
        handler, mock_client = handler_with_mock
        run_id = uuid4()

        handler.on_chain_start(
            serialized=None,
            inputs={"messages": []},
            run_id=run_id,
            metadata={"langgraph_node": "agent"},
            tags=["seq:step:3"],
            name="route_question",
        )

        handler.on_chain_end(
            outputs="search",
            run_id=run_id,
        )

        # Span should have been reclassified and decision attribute set
        if str(run_id) in handler._runs:
            span = handler._runs[str(run_id)]["span"]
            assert span.span_type == SpanType.CONDITIONAL

    def test_dict_output_stays_chain(self, handler_with_mock):
        """Chain with seq:step tag but dict output stays as CHAIN (not a router)."""
        handler, mock_client = handler_with_mock
        run_id = uuid4()

        handler.on_chain_start(
            serialized=None,
            inputs={"messages": []},
            run_id=run_id,
            metadata={"langgraph_node": "agent"},
            tags=["seq:step:3"],
            name="process_state",
        )

        assert str(run_id) in handler._potential_conditional_run_ids

        handler.on_chain_end(
            outputs={"messages": [{"role": "assistant", "content": "hi"}]},
            run_id=run_id,
        )

        # Should have been removed from potential set but NOT reclassified
        assert str(run_id) not in handler._potential_conditional_run_ids
        if str(run_id) in handler._runs:
            span = handler._runs[str(run_id)]["span"]
            assert span.span_type == SpanType.CHAIN

    def test_internal_names_excluded(self, handler_with_mock):
        """Known LangGraph internal names are not flagged as potential conditional edges."""
        handler, mock_client = handler_with_mock

        for internal_name in ["ChannelWrite", "ChannelRead", "RunnableSeq", "__start__", "__end__", "LangGraph"]:
            run_id = uuid4()
            handler.on_chain_start(
                serialized=None,
                inputs={},
                run_id=run_id,
                metadata={"langgraph_node": "agent"},
                tags=["seq:step:3"],
                name=internal_name,
            )
            assert str(run_id) not in handler._potential_conditional_run_ids

    def test_graph_step_tag_excluded(self, handler_with_mock):
        """Chains with graph:step tags are regular nodes, not conditional edges."""
        handler, mock_client = handler_with_mock
        run_id = uuid4()

        handler.on_chain_start(
            serialized=None,
            inputs={"query": "test"},
            run_id=run_id,
            metadata={"langgraph_node": "agent"},
            tags=["graph:step:1"],
            name="agent_node",
        )

        assert str(run_id) not in handler._potential_conditional_run_ids

    def test_no_langgraph_metadata_excluded(self, handler_with_mock):
        """Chains without LangGraph metadata are not flagged even with seq:step tags."""
        handler, mock_client = handler_with_mock
        run_id = uuid4()

        handler.on_chain_start(
            serialized=None,
            inputs={"query": "test"},
            run_id=run_id,
            metadata={"some_other_key": "value"},
            tags=["seq:step:1"],
            name="my_function",
        )

        assert str(run_id) not in handler._potential_conditional_run_ids

    def test_both_seq_and_graph_step_excluded(self, handler_with_mock):
        """Chains with both seq:step and graph:step tags are not flagged."""
        handler, mock_client = handler_with_mock
        run_id = uuid4()

        handler.on_chain_start(
            serialized=None,
            inputs={},
            run_id=run_id,
            metadata={"langgraph_node": "agent"},
            tags=["graph:step:1", "seq:step:1"],
            name="some_function",
        )

        assert str(run_id) not in handler._potential_conditional_run_ids

    def test_no_tags_excluded(self, handler_with_mock):
        """Chains without any tags are not flagged."""
        handler, mock_client = handler_with_mock
        run_id = uuid4()

        handler.on_chain_start(
            serialized=None,
            inputs={},
            run_id=run_id,
            metadata={"langgraph_node": "agent"},
            name="some_function",
        )

        assert str(run_id) not in handler._potential_conditional_run_ids


class TestHandoffSpanParenting:
    """Tests for HANDOFF span parent_id assignment."""

    @pytest.fixture
    def handler_with_mock(self):
        """Create a handler with mock client and simulated active trace."""
        with patch("lumenova_beacon.tracing.integrations.langchain.get_client") as mock_get:
            mock_client = MagicMock()
            mock_client.config.session_id = None
            mock_client.config.environment = None
            mock_get.return_value = mock_client
            handler = BeaconLangGraphHandler()
            yield handler, mock_client

    def test_handoff_parented_under_active_tool_span(self, handler_with_mock):
        """When a tool span is active, HANDOFF span uses it as parent."""
        handler, mock_client = handler_with_mock
        tool_run_id = uuid4()

        # Start a root chain to establish trace context
        root_run_id = uuid4()
        handler.on_chain_start(
            serialized={"name": "RootChain"},
            inputs={"query": "test"},
            run_id=root_run_id,
        )

        # Start a tool span
        handler.on_tool_start(
            serialized={"name": "TransferTool"},
            input_str="transfer to math agent",
            run_id=tool_run_id,
            parent_run_id=root_run_id,
        )
        tool_span = handler._runs[str(tool_run_id)]["span"]

        # Create handoff context — should parent under active tool span
        ctx = handler.handoff_context("Math Agent")
        with ctx:
            assert ctx.handoff_span is not None
            assert ctx.handoff_span.parent_id == tool_span.span_id

    def test_handoff_falls_back_to_root_span_when_no_tool_active(self, handler_with_mock):
        """When no tool span is active, HANDOFF span falls back to root span."""
        handler, mock_client = handler_with_mock

        # Start a root chain to establish trace context
        root_run_id = uuid4()
        handler.on_chain_start(
            serialized={"name": "RootChain"},
            inputs={"query": "test"},
            run_id=root_run_id,
        )
        root_span_id = handler._root_span_id
        assert root_span_id is not None

        # No tool is active — _active_tool_run_ids is empty
        assert handler.current_tool_span_id is None

        # Create handoff context — should fall back to root span
        ctx = handler.handoff_context("Math Agent")
        with ctx:
            assert ctx.handoff_span is not None
            assert ctx.handoff_span.parent_id == root_span_id

    def test_handoff_never_has_none_parent_when_root_exists(self, handler_with_mock):
        """HANDOFF span parent_id is never None when a root span exists."""
        handler, mock_client = handler_with_mock

        # Start a root chain
        root_run_id = uuid4()
        handler.on_chain_start(
            serialized={"name": "RootChain"},
            inputs={"query": "test"},
            run_id=root_run_id,
        )

        # Ensure no tools are active
        handler._active_tool_run_ids = []

        ctx = handler.handoff_context("Agent B")
        with ctx:
            assert ctx.handoff_span.parent_id is not None

    def test_handoff_prefers_tool_span_over_root(self, handler_with_mock):
        """When both tool and root exist, HANDOFF prefers the tool span."""
        handler, mock_client = handler_with_mock

        # Start root chain
        root_run_id = uuid4()
        handler.on_chain_start(
            serialized={"name": "RootChain"},
            inputs={"query": "test"},
            run_id=root_run_id,
        )
        root_span_id = handler._root_span_id

        # Start tool span
        tool_run_id = uuid4()
        handler.on_tool_start(
            serialized={"name": "TransferTool"},
            input_str="handoff",
            run_id=tool_run_id,
            parent_run_id=root_run_id,
        )
        tool_span_id = handler._runs[str(tool_run_id)]["span"].span_id

        ctx = handler.handoff_context("Agent C")
        with ctx:
            # Should be the tool span, NOT the root span
            assert ctx.handoff_span.parent_id == tool_span_id
            assert ctx.handoff_span.parent_id != root_span_id


class TestHandlerReuseReset:
    """Test that reusing a handler across multiple invocations creates separate traces."""

    @pytest.fixture
    def reusable_handler(self):
        """Create a handler suitable for reuse tests."""
        with patch("lumenova_beacon.integrations.langchain.get_client") as mock_get:
            mock_client = MagicMock(spec=BeaconClient)
            mock_client.should_sample = MagicMock(return_value=True)
            mock_client._timestamp_override = None
            mock_get.return_value = mock_client
            yield BeaconLangGraphHandler()

    def _simulate_invocation(self, handler):
        """Simulate a complete root-level invocation (start + end)."""
        run_id = uuid4()
        handler.on_chain_start(
            serialized={"name": "LangGraph", "id": ["langchain", "graph"]},
            inputs={"messages": [{"content": "hello"}]},
            run_id=run_id,
        )
        trace_id = handler._current_trace_id
        root_span_id = handler._root_span_id
        handler.on_chain_end(outputs={"result": "ok"}, run_id=run_id)
        return trace_id, root_span_id

    def test_reuse_handler_gets_new_trace_id(self, reusable_handler):
        """Second invocation on same handler should get a different trace_id."""
        trace_id_1, _ = self._simulate_invocation(reusable_handler)
        trace_id_2, _ = self._simulate_invocation(reusable_handler)

        assert trace_id_1 is not None
        assert trace_id_2 is not None
        assert trace_id_1 != trace_id_2

    def test_reuse_handler_gets_new_root_span_id(self, reusable_handler):
        """Second invocation on same handler should get a different root_span_id."""
        _, root_span_id_1 = self._simulate_invocation(reusable_handler)
        _, root_span_id_2 = self._simulate_invocation(reusable_handler)

        assert root_span_id_1 is not None
        assert root_span_id_2 is not None
        assert root_span_id_1 != root_span_id_2

    def test_reuse_handler_clears_runs(self, reusable_handler):
        """Runs dict should be empty at the start of the second invocation."""
        self._simulate_invocation(reusable_handler)
        # After first invocation, runs should have been populated then cleared on end
        # But _runs may still have entries depending on cleanup. The key test is that
        # second invocation resets state.
        run_id_2 = uuid4()
        reusable_handler.on_chain_start(
            serialized={"name": "LangGraph", "id": ["langchain", "graph"]},
            inputs={"messages": [{"content": "hello"}]},
            run_id=run_id_2,
        )
        # Only the new run should be in _runs (old ones were cleared by reset)
        assert str(run_id_2) in reusable_handler._runs
        assert len(reusable_handler._runs) == 1

    def test_no_reset_on_resume(self, reusable_handler):
        """When _is_resume=True, trace_id and root_span_id must be preserved."""
        trace_id_1, root_span_id_1 = self._simulate_invocation(reusable_handler)

        # Set resume mode with the original trace/root IDs
        reusable_handler._is_resume = True
        reusable_handler._current_trace_id = trace_id_1
        reusable_handler._root_span_id = root_span_id_1

        # Start a new root chain - should NOT reset because is_resume=True
        run_id = uuid4()
        reusable_handler.on_chain_start(
            serialized={"name": "LangGraph", "id": ["langchain", "graph"]},
            inputs={"messages": [{"content": "resumed"}]},
            run_id=run_id,
        )

        assert reusable_handler._current_trace_id == trace_id_1
        assert reusable_handler._root_span_id == root_span_id_1

    def test_no_reset_on_first_invocation(self, reusable_handler):
        """First invocation should not trigger a reset (no prior completed invocation)."""
        # Pre-set trace_id (simulating BeaconLangGraphConfig behavior)
        preset_trace_id = "aabbccdd" * 4
        preset_root_span_id = "11223344" * 2
        reusable_handler._current_trace_id = preset_trace_id
        reusable_handler._root_span_id = preset_root_span_id

        # First invocation - _invocation_completed is False, so no reset should happen
        run_id = uuid4()
        reusable_handler.on_chain_start(
            serialized={"name": "LangGraph", "id": ["langchain", "graph"]},
            inputs={"messages": [{"content": "hello"}]},
            run_id=run_id,
        )

        # Pre-set trace_id should be used (not overwritten by reset)
        span = reusable_handler._runs[str(run_id)]["span"]
        assert span.trace_id == preset_trace_id

    def test_no_reset_for_non_root_span(self, reusable_handler):
        """Non-root spans (with parent_run_id) should never trigger a reset."""
        trace_id_1, _ = self._simulate_invocation(reusable_handler)

        # Start a new root span for the second invocation
        root_run_id = uuid4()
        reusable_handler.on_chain_start(
            serialized={"name": "LangGraph", "id": ["langchain", "graph"]},
            inputs={"messages": [{"content": "hello"}]},
            run_id=root_run_id,
        )
        trace_id_2 = reusable_handler._current_trace_id

        # Now start a child span - should NOT reset
        child_run_id = uuid4()
        reusable_handler.on_chain_start(
            serialized={"name": "agent", "id": ["langchain", "chain"]},
            inputs={},
            run_id=child_run_id,
            parent_run_id=root_run_id,
        )

        # trace_id should still be the second invocation's trace_id
        assert reusable_handler._current_trace_id == trace_id_2
        assert reusable_handler._current_trace_id != trace_id_1
