"""LangChain integration using SDK Spans.

This module provides a callback handler that traces LangChain operations
using SDK Span objects, which are exported via OTLP.

Usage:
    from lumenova_beacon import BeaconClient
    from lumenova_beacon import BeaconLangGraphHandler
    from langchain_openai import ChatOpenAI

    # Initialize BeaconClient (sets up TracerProvider)
    client = BeaconClient()

    # Create callback handler
    handler = BeaconLangGraphHandler()

    # Use with LangChain
    llm = ChatOpenAI()
    response = llm.invoke("Hello", config={"callbacks": [handler]})

Environment Variables:
    BEACON_ENDPOINT - Beacon API endpoint (for OTLP export)
    BEACON_API_KEY - Beacon API key (for OTLP export)
    BEACON_SESSION_ID - Default session ID (optional)
"""

from __future__ import annotations

import json
import logging
import time
from contextvars import ContextVar, Token
from dataclasses import dataclass, field
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Sequence
from uuid import UUID

from lumenova_beacon.core.client import get_client
from lumenova_beacon.tracing.span import Span
from lumenova_beacon.tracing.trace import (
    get_current_span,
    set_current_span,
    clear_context,
)
from lumenova_beacon.types import SpanType, SpanKind, StatusCode
from lumenova_beacon.utils import generate_trace_id, generate_span_id

if TYPE_CHECKING:
    from lumenova_beacon.governance.config import GovernanceConfig

try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.agents import AgentAction, AgentFinish
    from langchain_core.documents import Document
    from langchain_core.messages import BaseMessage
    from langchain_core.outputs import LLMResult
except ImportError:
    raise ImportError(
        "langchain-core is required for LangChain OTLP integration. "
        "Install it with: pip install langchain-core"
    )

logger = logging.getLogger(__name__)

# Keys used for LangGraph checkpoint persistence. Filtered out when writing
# metadata as span attributes to avoid duplicating data already on the span.
_BEACON_CHECKPOINT_KEYS = frozenset({
    "beacon_trace_id",
    "beacon_root_span_id",
    "beacon_session_id",
})


@dataclass
class _HandoffState:
    """Isolated per-handoff working state.

    Stored in a ContextVar so each async task gets its own copy,
    enabling parallel handoffs on the same handler instance.
    """

    trace_id: str | None
    root_span_id: str | None
    runs: dict[str, dict[str, Any]]
    active_tool_run_ids: list[str]
    langgraph_parent_ids: set[str]
    first_token_times: dict[str, float | None]
    skipped_resume_roots: dict[str, dict[str, Any]]
    is_resume: bool
    parent_handoff_trace_id: str | None
    parent_handoff_span_id: str | None
    parent_handoff_agent_name: str | None
    agent_name: str | None
    interrupt_detected: bool
    interrupt_info: dict[str, Any]
    potential_conditional_run_ids: set[str] = field(default_factory=set)
    root_span_created: bool = False


# Task-local handoff state. When set, callbacks use this state instead
# of the handler's base instance attributes. Each async task gets its
# own copy automatically, enabling parallel handoffs.
_active_handoff_state: ContextVar[_HandoffState | None] = ContextVar(
    "_active_handoff_state", default=None
)


class _StateAccessor:
    """Uniform read/write access to handler state.

    Delegates to the active ``_HandoffState`` (from ContextVar) when
    inside a handoff, otherwise to the handler's base attributes.
    This lets callback methods work identically regardless of whether
    a handoff is active, and enables parallel handoffs by isolating
    mutable state per async task.
    """

    __slots__ = ("_hs", "_h")

    def __init__(self, handler: "BeaconLangGraphHandler"):
        self._hs: _HandoffState | None = _active_handoff_state.get()
        self._h: BeaconLangGraphHandler = handler

    @property
    def in_handoff(self) -> bool:
        return self._hs is not None

    # -- trace / root --

    @property
    def current_trace_id(self) -> str | None:
        return self._hs.trace_id if self._hs else self._h._current_trace_id

    @current_trace_id.setter
    def current_trace_id(self, value: str | None) -> None:
        if self._hs:
            self._hs.trace_id = value
        else:
            self._h._current_trace_id = value

    @property
    def root_span_id(self) -> str | None:
        return self._hs.root_span_id if self._hs else self._h._root_span_id

    @root_span_id.setter
    def root_span_id(self, value: str | None) -> None:
        if self._hs:
            self._hs.root_span_id = value
        else:
            self._h._root_span_id = value

    # -- runs --

    @property
    def runs(self) -> dict[str, dict[str, Any]]:
        return self._hs.runs if self._hs else self._h._runs

    @runs.setter
    def runs(self, value: dict[str, dict[str, Any]]) -> None:
        if self._hs:
            self._hs.runs = value
        else:
            self._h._runs = value

    # -- tool tracking --

    @property
    def active_tool_run_ids(self) -> list[str]:
        return self._hs.active_tool_run_ids if self._hs else self._h._active_tool_run_ids

    @active_tool_run_ids.setter
    def active_tool_run_ids(self, value: list[str]) -> None:
        if self._hs:
            self._hs.active_tool_run_ids = value
        else:
            self._h._active_tool_run_ids = value

    # -- langgraph parents --

    @property
    def langgraph_parent_ids(self) -> set[str]:
        return self._hs.langgraph_parent_ids if self._hs else self._h._langgraph_parent_ids

    # -- TTFT --

    @property
    def first_token_times(self) -> dict[str, float | None]:
        return self._hs.first_token_times if self._hs else self._h._first_token_times

    # -- resume --

    @property
    def skipped_resume_roots(self) -> dict[str, dict[str, Any]]:
        return self._hs.skipped_resume_roots if self._hs else self._h._skipped_resume_roots

    @property
    def is_resume(self) -> bool:
        return self._hs.is_resume if self._hs else self._h._is_resume

    # -- handoff back-link --

    @property
    def parent_handoff_trace_id(self) -> str | None:
        return self._hs.parent_handoff_trace_id if self._hs else self._h._parent_handoff_trace_id

    @property
    def parent_handoff_span_id(self) -> str | None:
        return self._hs.parent_handoff_span_id if self._hs else self._h._parent_handoff_span_id

    @property
    def parent_handoff_agent_name(self) -> str | None:
        return self._hs.parent_handoff_agent_name if self._hs else self._h._parent_handoff_agent_name

    # -- agent name --

    @property
    def agent_name(self) -> str | None:
        return self._hs.agent_name if self._hs else self._h._agent_name

    @agent_name.setter
    def agent_name(self, value: str | None) -> None:
        if self._hs:
            self._hs.agent_name = value
        else:
            self._h._agent_name = value

    # -- interrupt --

    @property
    def interrupt_detected(self) -> bool:
        return self._hs.interrupt_detected if self._hs else self._h._interrupt_detected

    @interrupt_detected.setter
    def interrupt_detected(self, value: bool) -> None:
        if self._hs:
            self._hs.interrupt_detected = value
        else:
            self._h._interrupt_detected = value

    @property
    def interrupt_info(self) -> dict[str, Any]:
        return self._hs.interrupt_info if self._hs else self._h._interrupt_info

    @interrupt_info.setter
    def interrupt_info(self, value: dict[str, Any]) -> None:
        if self._hs:
            self._hs.interrupt_info = value
        else:
            self._h._interrupt_info = value

    # -- conditional routing --

    @property
    def potential_conditional_run_ids(self) -> set[str]:
        return self._hs.potential_conditional_run_ids if self._hs else self._h._potential_conditional_run_ids

    # -- current_tool_span_id (derived) --

    @property
    def current_tool_span_id(self) -> str | None:
        for tool_run_id in reversed(self.active_tool_run_ids):
            if tool_run_id in self.runs:
                return self.runs[tool_run_id]["span"].span_id
        return None


class HandoffContext:
    """Context manager for handler-level agent handoff.

    Creates a HANDOFF span in the parent trace and isolates the child
    agent's callback state using a ``ContextVar``.  Each async task gets
    its own copy of the ContextVar automatically, so parallel handoffs
    on the same handler instance never interfere with each other.

    The handler's base instance attributes are **never mutated** during
    a handoff — all child state lives in the ContextVar-scoped
    ``_HandoffState``.

    Usage::

        async with handler.handoff_context("Math Agent") as ctx:
            result = await agent.graph.ainvoke(
                {"messages": [...]},
            )

    Attributes:
        trace_id: Pre-generated trace_id for the child agent.
        root_span_id: Pre-generated span_id for the child's root span.
        handoff_span: The HANDOFF span in the parent trace.
    """

    def __init__(
        self,
        handler: "BeaconLangGraphHandler",
        target_agent: str,
        config: dict[str, Any] | None = None,
        parent_tool_run_id: str | None = None,
    ):
        self._handler = handler
        self._target_agent = target_agent
        self._config = config
        self._parent_tool_run_id = parent_tool_run_id
        self._cv_token: Token[_HandoffState | None] | None = None

        # Pre-generate child trace identifiers (may be overridden in _enter
        # if auto-resume detects a previous interrupted handoff)
        self.trace_id = generate_trace_id()
        self.root_span_id = generate_span_id()
        self._is_resume = False
        self.handoff_span: Span | None = None

    def _resolve_tool_parent(self, s: _StateAccessor) -> str | None:
        """Try to find the span_id of the tool that owns this handoff.

        When ``parent_tool_run_id`` is provided (from auto-detected
        handoffs), we directly look up the tool span — this is the most
        precise method since the subgraph root's parent_run_id identifies
        the exact tool.

        When ``config`` is provided (from a LangChain tool body), the
        callback manager's ``parent_run_id`` identifies the tools chain
        node that launched the parallel tool invocations. We use this to
        find an unclaimed active tool span that is a child of that node,
        ensuring each parallel handoff parents under a different tool.

        Falls back to ``current_tool_span_id`` (last active tool) when
        neither is available.
        """
        # Direct tool run_id (auto-detected handoffs know the exact tool)
        if self._parent_tool_run_id and self._parent_tool_run_id in s.runs:
            tool_span = s.runs[self._parent_tool_run_id]["span"]
            return tool_span.span_id

        if self._config:
            callbacks = self._config.get("callbacks")
            parent_run_id = getattr(callbacks, "parent_run_id", None) if callbacks else None
            if parent_run_id:
                parent_key = str(parent_run_id)

                # parent_run_id may be the tool's own run_id (single tool)
                # or the tools chain node (parallel tools). Check both.
                if parent_key in s.runs:
                    parent_span = s.runs[parent_key]["span"]

                    # Check if it's directly a tool in active_tool_run_ids
                    if parent_key in s.active_tool_run_ids:
                        return parent_span.span_id

                    # Otherwise it's the tools chain node — claim an
                    # unclaimed child tool span.
                    parent_span_id = parent_span.span_id
                    handler = self._handler
                    for tool_run_id in s.active_tool_run_ids:
                        if tool_run_id in handler._claimed_tool_run_ids:
                            continue
                        if tool_run_id in s.runs:
                            tool_span = s.runs[tool_run_id]["span"]
                            if tool_span.parent_id == parent_span_id:
                                handler._claimed_tool_run_ids.add(tool_run_id)
                                return tool_span.span_id

        # Fallback: last active tool (works for single non-parallel tool)
        return s.current_tool_span_id

    def _enter(self) -> "HandoffContext":
        handler = self._handler
        # Read current effective state (supports nested handoffs)
        s = _StateAccessor(handler)

        # 0. Auto-detect resume: if a previous handoff to the same agent was
        #    interrupted within this parent trace, reuse its child trace IDs.
        if not self._is_resume and s.current_trace_id:
            key = (s.current_trace_id, self._target_agent)
            saved = handler._interrupted_handoff_traces.get(key)
            if saved:
                self.trace_id, self.root_span_id = saved
                self._is_resume = True

        # 1. Determine parent for the handoff span (active tool span, or root).
        #    When config is provided (from a LangChain tool body), we can
        #    resolve the exact tool run_id via the callback manager's
        #    parent_run_id — essential for parallel tool invocations.
        parent_span_id = self._resolve_tool_parent(s) or s.root_span_id

        # 2. Create HANDOFF span in the PARENT trace
        self.handoff_span = Span(
            name=f"handoff:{self._target_agent}",
            trace_id=s.current_trace_id,
            parent_id=parent_span_id,
            span_type=SpanType.HANDOFF,
            session_id=handler._session_id,
            user_email=handler._user_email,
            user_name=handler._user_name,
            user_id=handler._user_id,
            record_stacktrace=handler._record_stacktrace,
        )
        self.handoff_span.set_attribute("handoff.target_agent", self._target_agent)
        self.handoff_span.set_attribute("handoff.target_trace_id", self.trace_id)
        self.handoff_span.set_attribute("handoff.target_root_span_id", self.root_span_id)
        if self._is_resume:
            self.handoff_span.set_attribute("handoff.is_resume", True)
        self.handoff_span.add_link(
            trace_id=self.trace_id,
            span_id=self.root_span_id,
            attributes={"link.type": "child_trace"},
        )
        self.handoff_span.start()

        # 3. Create isolated child state in ContextVar (task-local)
        child_state = _HandoffState(
            trace_id=self.trace_id,
            root_span_id=self.root_span_id,
            runs={},
            active_tool_run_ids=[],
            langgraph_parent_ids=set(),
            first_token_times={},
            skipped_resume_roots={},
            is_resume=self._is_resume,
            parent_handoff_trace_id=s.current_trace_id,
            parent_handoff_span_id=self.handoff_span.span_id,
            parent_handoff_agent_name=s.agent_name,
            agent_name=self._target_agent,
            interrupt_detected=False,
            interrupt_info={},
            potential_conditional_run_ids=set(),
            # On resume the root span already exists from the first handoff;
            # mark it so the AGENT-type override doesn't fire again.
            root_span_created=self._is_resume,
        )
        self._cv_token = _active_handoff_state.set(child_state)

        return self

    def _exit(self, exc_val: BaseException | None) -> None:
        handler = self._handler
        assert self.handoff_span is not None
        assert self._cv_token is not None

        # 1. Flush orphaned child spans BEFORE restoring ContextVar.
        #    The ContextVar still points to the child state, so
        #    _flush_orphaned_spans will find in-progress child LLM spans.
        if exc_val is not None and (
            handler._is_cancelled_error(exc_val)
            or handler._is_langgraph_interrupt(exc_val)
        ):
            handler._flush_orphaned_spans()

        # 2. End the handoff span with proper status
        if exc_val is not None:
            if handler._is_langgraph_interrupt(exc_val):
                # Interrupt is not an error — normal HITL flow
                interrupt_info = handler._extract_interrupt_info(exc_val)
                self.handoff_span.set_attribute(handler.ATTR_LANGGRAPH_INTERRUPT, True)
                if interrupt_info.get("id"):
                    self.handoff_span.set_attribute(
                        handler.ATTR_LANGGRAPH_INTERRUPT_ID, interrupt_info["id"]
                    )
                if interrupt_info.get("value"):
                    self.handoff_span.set_attribute(
                        handler.ATTR_LANGGRAPH_INTERRUPT_VALUE,
                        str(interrupt_info["value"]),
                    )
                self.handoff_span.set_status(
                    StatusCode.OK,
                    f"Interrupted: {interrupt_info.get('value', 'awaiting input')}",
                )
                # Save child trace IDs so the next handoff to the same agent
                # within this parent trace auto-resumes the child trace.
                parent_trace_id = self.handoff_span.trace_id
                if parent_trace_id:
                    key = (parent_trace_id, self._target_agent)
                    handler._interrupted_handoff_traces[key] = (
                        self.trace_id,
                        self.root_span_id,
                    )
            elif handler._is_cancelled_error(exc_val):
                self.handoff_span.set_attribute(handler.ATTR_LANGGRAPH_CANCELLED, True)
                self.handoff_span.set_status(StatusCode.OK, "Cancelled by user")
            else:
                self.handoff_span.set_attribute("error.type", type(exc_val).__qualname__)
                self.handoff_span.record_exception(exc_val)
        else:
            if self.handoff_span.status_code == StatusCode.UNSET:
                self.handoff_span.set_status(StatusCode.OK)
            # Normal completion — clean up any saved resume state
            parent_trace_id = self.handoff_span.trace_id
            if parent_trace_id:
                key = (parent_trace_id, self._target_agent)
                handler._interrupted_handoff_traces.pop(key, None)

        if self.handoff_span.end_time is None:
            self.handoff_span.end()

        handler.client.export_span(self.handoff_span)

        # 3. Restore ContextVar to previous value (parent or None)
        try:
            _active_handoff_state.reset(self._cv_token)
        except ValueError:
            # Token was created in a different async Context (happens with
            # auto-detected handoffs where on_chain_start and on_chain_end
            # fire in different executor threads). Fall back to set(None).
            _active_handoff_state.set(None)
        self._cv_token = None

    def __enter__(self) -> "HandoffContext":
        return self._enter()

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self._exit(exc_val)

    async def __aenter__(self) -> "HandoffContext":
        return self._enter()

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        self._exit(exc_val)


def get_beacon_handler(config: dict[str, Any]) -> "BeaconLangGraphHandler | None":
    """Extract the BeaconLangGraphHandler from a LangChain RunnableConfig.

    Works with both raw callback lists and CallbackManager instances.

    Args:
        config: LangChain RunnableConfig dict

    Returns:
        The BeaconLangGraphHandler if found, None otherwise
    """
    callbacks = config.get("callbacks") or []
    # Handle CallbackManager (has .handlers) vs plain list
    if hasattr(callbacks, "handlers"):
        callbacks = callbacks.handlers
    for cb in callbacks:
        if isinstance(cb, BeaconLangGraphHandler):
            return cb
    return None


class BeaconLangGraphHandler(BaseCallbackHandler):
    """Callback handler for tracing LangChain operations.

    Traces LangChain operations using SDK Span objects, which are exported
    via OTLP to the configured backend.

    Handles:
    - Chain operations (on_chain_start/end/error)
    - LLM calls (on_llm_start/end/error, on_chat_model_start)
    - Tool invocations (on_tool_start/end/error)
    - Retriever queries (on_retriever_start/end/error)
    - Agent actions (on_agent_action/finish)

    Example:
        >>> from lumenova_beacon import BeaconClient
        >>> from lumenova_beacon import BeaconLangGraphHandler
        >>> from langchain_openai import ChatOpenAI
        >>>
        >>> # Initialize BeaconClient (sets up TracerProvider)
        >>> client = BeaconClient()
        >>>
        >>> # Create callback handler
        >>> handler = BeaconLangGraphHandler()
        >>>
        >>> # Use with LangChain
        >>> llm = ChatOpenAI()
        >>> response = llm.invoke("Hello", config={"callbacks": [handler]})
    """

    # Span type attributes
    ATTR_SPAN_TYPE = "beacon.span_type"
    ATTR_COMPONENT_TYPE = "langchain.component_type"

    # GenAI semantic conventions
    # See: https://opentelemetry.io/docs/specs/semconv/gen-ai/
    ATTR_LLM_MODEL = "gen_ai.request.model"

    # OTEL GenAI semantic conventions
    ATTR_LLM_INPUT_MESSAGES = "gen_ai.input.messages"
    ATTR_LLM_OUTPUT_MESSAGES = "gen_ai.output.messages"
    ATTR_LLM_PROVIDER_NAME = "gen_ai.provider.name"
    ATTR_LLM_INPUT_TOKENS = "gen_ai.usage.input_tokens"
    ATTR_LLM_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
    ATTR_LLM_REASONING_TOKENS = "gen_ai.usage.reasoning_tokens"
    ATTR_LLM_CACHE_READ_TOKENS = "gen_ai.usage.cache_read_input_tokens"
    ATTR_LLM_CACHE_CREATION_TOKENS = "gen_ai.usage.cache_creation_input_tokens"
    ATTR_OPERATION_NAME = "gen_ai.operation.name"
    ATTR_RESPONSE_MODEL = "gen_ai.response.model"
    ATTR_RESPONSE_ID = "gen_ai.response.id"
    ATTR_RESPONSE_FINISH_REASONS = "gen_ai.response.finish_reasons"
    ATTR_REQUEST_STOP_SEQUENCES = "gen_ai.request.stop_sequences"
    ATTR_REQUEST_SEED = "gen_ai.request.seed"
    ATTR_TOOL_DEFINITIONS = "gen_ai.tool.definitions"

    # OTEL GenAI tool span attributes
    ATTR_TOOL_NAME = "gen_ai.tool.name"
    ATTR_TOOL_TYPE = "gen_ai.tool.type"
    ATTR_TOOL_CALL_ARGUMENTS = "gen_ai.tool.call.arguments"
    ATTR_TOOL_CALL_RESULT = "gen_ai.tool.call.result"

    # OTEL GenAI conversation attribute
    ATTR_CONVERSATION_ID = "gen_ai.conversation.id"

    # TTFT (Time-to-First-Token) attributes
    ATTR_TTFT_MS = "gen_ai.response.time_to_first_token_ms"
    ATTR_STREAMING = "gen_ai.request.streaming"

    # Class-level storage for interrupt IDs keyed by trace_id.
    # Allows resume handlers to automatically pick up the interrupt ID
    # from the previous invocation on the same trace.
    _trace_interrupt_ids: dict[str, str] = {}

    # Class-level storage for interrupted handoff child trace IDs.
    # When a handoff exits with an interrupt, the child trace IDs are saved
    # keyed by (parent_trace_id, target_agent). On the next handoff to the
    # same agent within the same parent trace, the child trace is resumed
    # automatically instead of generating fresh IDs.
    _interrupted_handoff_traces: dict[tuple[str, str], tuple[str, str]] = {}

    # GenAI Agent attributes (OTEL standard)
    ATTR_AGENT_NAME = "gen_ai.agent.name"
    ATTR_AGENT_FRAMEWORK = "gen_ai.agent.framework"

    # Deployment attributes (OTEL standard)
    ATTR_DEPLOYMENT_ENVIRONMENT = "deployment.environment.name"

    # Model request parameters (OTEL standard)
    ATTR_REQUEST_TEMPERATURE = "gen_ai.request.temperature"
    ATTR_REQUEST_TOP_P = "gen_ai.request.top_p"
    ATTR_REQUEST_TOP_K = "gen_ai.request.top_k"
    ATTR_REQUEST_MAX_TOKENS = "gen_ai.request.max_tokens"
    ATTR_REQUEST_FREQUENCY_PENALTY = "gen_ai.request.frequency_penalty"
    ATTR_REQUEST_PRESENCE_PENALTY = "gen_ai.request.presence_penalty"

    # Beacon-specific attributes
    ATTR_BEACON_SESSION_ID = "beacon.session_id"
    ATTR_BEACON_METADATA_PREFIX = "beacon.metadata"

    # LangGraph-specific attributes
    ATTR_LANGGRAPH_INTERRUPT = "langgraph.interrupt"
    ATTR_LANGGRAPH_INTERRUPT_ID = "langgraph.interrupt.id"
    ATTR_LANGGRAPH_INTERRUPT_VALUE = "langgraph.interrupt.value"
    ATTR_LANGGRAPH_THREAD_ID = "langgraph.thread_id"
    ATTR_LANGGRAPH_RESUME = "langgraph.resume"
    ATTR_LANGGRAPH_RESUME_VALUE = "langgraph.resume.value"
    ATTR_LANGGRAPH_CANCELLED = "langgraph.cancelled"
    ATTR_LANGGRAPH_ROUTING_DECISION = "langgraph.routing_decision"

    # Known LangGraph internal chain names (not conditional routing functions)
    _LANGGRAPH_INTERNAL_NAMES = frozenset({
        "ChannelWrite",
        "ChannelRead",
        "RunnableSeq",
        "__start__",
        "__end__",
        "LangGraph",
    })

    def __init__(
        self,
        session_id: str | None = None,
        user_email: str | None = None,
        user_name: str | None = None,
        user_id: str | None = None,
        environment: str | None = None,
        agent_name: str | None = None,
        metadata: dict[str, Any] | None = None,
        autodetect_handoffs: bool = False,
        record_stacktrace: bool | None = None,
        **kwargs: Any,
    ):
        """Initialize the LangChain callback handler.

        Args:
            session_id: Default session ID for all spans (optional)
            user_email: User email for attribution (optional, falls back to client config)
            user_name: User display name for attribution (optional, falls back to client config)
            user_id: User ID for attribution (optional, falls back to client config)
            environment: Environment name (e.g., "production", "staging", "development")
            agent_name: Human-readable name of the GenAI agent
            metadata: Custom metadata dict to add to all spans (key-value pairs)
            autodetect_handoffs: When True, nested graph invocations
                under tools that have LangGraph metadata are automatically
                treated as handoffs (separate traces with HANDOFF spans).
            record_stacktrace: Capture full stack traces on exceptions (default: inherits from client config)
        """
        super().__init__()

        # Cache the client instance to avoid creating duplicates in async contexts
        self._client = get_client()

        self._session_id = session_id if session_id is not None else self._client.config.session_id
        self._user_email = user_email if user_email is not None else self._client.config.user_email
        self._user_name = user_name if user_name is not None else self._client.config.user_name
        self._user_id = user_id if user_id is not None else self._client.config.user_id
        self._environment = environment if environment is not None else self._client.config.environment
        self._record_stacktrace = record_stacktrace if record_stacktrace is not None else self._client.config.record_stacktrace
        self._agent_name = agent_name
        self._metadata = metadata or {}
        self._metadata.update(kwargs)  # Allow additional metadata via kwargs

        # Track active spans by run_id
        self._runs: dict[str, dict[str, Any]] = {}
        self._langgraph_parent_ids: set[str] = set()

        # TTFT tracking for streaming calls
        self._first_token_times: dict[str, float | None] = {}

        # Trace continuation support for LangGraph interrupts
        self._current_trace_id: str | None = None  # Trace ID to use (or set when first span created)
        self._skipped_resume_roots: dict[str, dict[str, Any]] = {}  # run_id -> {parent_id, resume_value}
        self._root_span_id: str | None = None  # Root span ID (captured from first span, or set by BeaconLangGraphConfig)
        self._is_resume = False  # Whether this is a resume call
        self._interrupt_id: str | None = None  # Interrupt ID from previous invocation (for resume linking)

        # Interrupt propagation: tracks whether any child detected an interrupt
        # (root may receive on_chain_end instead of on_chain_error for interrupts)
        self._interrupt_detected: bool = False
        self._interrupt_info: dict[str, Any] = {}

        # Track active tools for nested graph parenting
        self._active_tool_run_ids: list[str] = []  # Stack of active tool run_ids
        self._claimed_tool_run_ids: set[str] = set()  # Tools already claimed by a handoff

        # Conditional edge detection: tracks run_ids that may be routing functions
        self._potential_conditional_run_ids: set[str] = set()

        # Tracks whether a full invocation has completed on this handler.
        # Used to detect handler reuse across multiple invoke() calls.
        self._invocation_completed: bool = False

        # Handoff back-link: if set, the root span gets an OTEL link to the parent handoff
        self._parent_handoff_trace_id: str | None = None
        self._parent_handoff_span_id: str | None = None
        self._parent_handoff_agent_name: str | None = None

        # Auto-detect subagent handoffs: when enabled, nested graph invocations
        # under tools are automatically treated as handoffs (separate traces).
        # _auto_handoff_contexts is NOT part of _HandoffState, so it persists
        # across state swaps — critical for nested auto-handoffs.
        self._autodetect_handoffs = autodetect_handoffs
        self._auto_handoff_contexts: dict[str, HandoffContext] = {}
        # Handler-level state for auto-handoff cross-context support.
        # ContextVar isolation doesn't work across async tasks (LangChain
        # fires callbacks in executor threads with copied Contexts), so we
        # keep a handler-level mapping and restore the ContextVar at the
        # top of each callback.
        self._auto_handoff_child_states: dict[str, _HandoffState] = {}  # root run_id -> child state
        self._auto_handoff_run_map: dict[str, str] = {}  # any run_id -> root run_id

        logger.debug("Initialized BeaconLangGraphHandler")

    def _reset_invocation_state(self) -> None:
        """Reset per-invocation state for handler reuse across multiple calls.

        Called at the start of a new top-level invocation (root on_chain_start)
        when the handler has leftover state from a previous invocation.
        Does NOT reset configuration attributes (session_id, environment,
        agent_name, metadata, autodetect_handoffs), externally-set resume
        state (_is_resume, _interrupt_id), or _auto_handoff_contexts.
        """
        self._current_trace_id = None
        self._root_span_id = None
        self._runs = {}
        self._langgraph_parent_ids = set()
        self._first_token_times = {}
        self._skipped_resume_roots = {}
        self._interrupt_detected = False
        self._interrupt_info = {}
        self._active_tool_run_ids = []
        self._claimed_tool_run_ids = set()
        self._potential_conditional_run_ids = set()
        self._invocation_completed = False
        self._auto_handoff_child_states = {}
        self._auto_handoff_run_map = {}

    @property
    def client(self):
        """Get the cached BeaconClient instance."""
        return self._client

    @property
    def trace_id(self) -> str | None:
        """Get the current trace ID.

        This is the trace_id of the root span, or the continuation trace_id
        if set. Use this to continue the trace in subsequent requests.

        Returns:
            The current trace_id, or None if no trace has been started
        """
        return self._current_trace_id

    @property
    def root_span_id(self) -> str | None:
        """Get the root span ID for resume continuation.

        Returns:
            The root span_id, or None if no trace has been started
        """
        return self._root_span_id

    @property
    def interrupt_id(self) -> str | None:
        """Get the interrupt checkpoint ID from the last detected interrupt.

        Returns:
            The interrupt checkpoint ID, or None if no interrupt was detected
        """
        return self._interrupt_info.get("id") or self._interrupt_id

    @property
    def interrupt_detected(self) -> bool:
        """Whether an interrupt was detected during execution."""
        return self._interrupt_detected

    def mark_interrupt(
        self,
        value: str | None = None,
        interrupt_id: str | None = None,
    ) -> None:
        """Mark the current trace as interrupted by a breakpoint.

        Use this when ``interrupt_before`` / ``interrupt_after`` breakpoints
        are detected externally (e.g. via ``__interrupt__`` in
        ``astream_events``).  These breakpoints do **not** raise exceptions
        through LangChain callbacks, so the handler cannot detect them on
        its own.

        If called **before** the root span's ``on_chain_end``, the existing
        propagation logic will set the interrupt attributes automatically.
        If called **after** the root span has already been exported, this
        method re-exports the root span with the interrupt attributes so the
        backend can merge them.

        Args:
            value: Human-readable interrupt prompt / description.
            interrupt_id: Checkpoint ID for linking the subsequent resume.
        """
        s = _StateAccessor(self)
        s.interrupt_detected = True
        s.interrupt_info = {
            "value": value or "Awaiting user input",
            "id": interrupt_id,
        }

        # Persist at class level so resume handlers can pick up the ID
        if interrupt_id and s.current_trace_id:
            BeaconLangGraphHandler._trace_interrupt_ids[s.current_trace_id] = interrupt_id

        # If the root span has already been exported (no longer in _runs),
        # re-export it with interrupt attributes.
        if s.root_span_id and s.current_trace_id:
            root_still_active = any(
                data["span"].span_id == s.root_span_id
                for data in s.runs.values()
            )
            if not root_still_active:
                self._reexport_root_with_interrupt()

    def _reexport_root_with_interrupt(self) -> None:
        """Re-export the root span with interrupt attributes.

        Used when an interrupt is detected after the root span has already
        been exported (e.g. ``mark_interrupt`` called after ``on_chain_end``).
        """
        s = _StateAccessor(self)
        if not s.root_span_id or not s.current_trace_id:
            return

        root_span = Span(
            name=s.agent_name or "agent",
            trace_id=s.current_trace_id,
            span_id=s.root_span_id,
            span_type=SpanType.AGENT,
            session_id=self._session_id,
            user_email=self._user_email,
            user_name=self._user_name,
            user_id=self._user_id,
        )
        root_span.set_attribute(self.ATTR_LANGGRAPH_INTERRUPT, True)
        if s.interrupt_info.get("id"):
            root_span.set_attribute(
                self.ATTR_LANGGRAPH_INTERRUPT_ID,
                s.interrupt_info["id"],
            )
        if s.interrupt_info.get("value"):
            root_span.set_attribute(
                self.ATTR_LANGGRAPH_INTERRUPT_VALUE,
                str(s.interrupt_info["value"]),
            )
        root_span.set_attribute(self.ATTR_AGENT_FRAMEWORK, "langgraph")
        if s.agent_name:
            root_span.set_attribute(self.ATTR_AGENT_NAME, s.agent_name)

        root_span.set_status(
            StatusCode.OK,
            f"Interrupted: {s.interrupt_info.get('value', 'awaiting input')}",
        )
        root_span.start()
        root_span.end()
        self.client.export_span(root_span)

        logger.debug(
            f"Re-exported root span with interrupt: trace_id={s.current_trace_id}, "
            f"span_id={s.root_span_id}, value={s.interrupt_info.get('value')}"
        )

    @property
    def current_tool_span_id(self) -> str | None:
        """Get the span ID of the currently active tool span.

        This is useful for placing handoff spans under the correct tool span
        when SDK context vars are not propagated (e.g., in async LangGraph
        tool execution).

        Returns:
            The active tool span_id, or None if no tool is active
        """
        return _StateAccessor(self).current_tool_span_id

    def handoff_context(
        self,
        target_agent: str,
        config: dict[str, Any] | None = None,
    ) -> HandoffContext:
        """Create a handoff context that isolates handler state for a sub-agent.

        Creates a HANDOFF span in the current (parent) trace, then sets a
        task-local ``ContextVar`` so all LangChain callbacks during the
        child's execution produce spans in a new independent trace.

        If a previous handoff to the same agent was interrupted, the child
        trace is automatically resumed (same trace_id and root_span_id).

        Usage::

            async with handler.handoff_context("Math Agent", config=config) as ctx:
                result = await agent.graph.ainvoke(
                    {"messages": [...]},
                    config={"configurable": {"thread_id": str(uuid.uuid4())}},
                )

        Args:
            target_agent: Display name of the target agent. Also used as the
                agent_name for child spans during the handoff.
            config: Optional LangChain ``RunnableConfig`` from the calling
                tool.  When provided, the HANDOFF span is correctly parented
                under that tool's span — essential for parallel tool
                invocations where multiple handoffs run concurrently.

        Returns:
            HandoffContext - supports both sync and async context manager.
        """
        return HandoffContext(
            handler=self,
            target_agent=target_agent,
            config=config,
        )

    def _auto_handoff_enter(
        self, run_id: UUID, subagent_name: str, tool_run_id: UUID | None = None
    ) -> None:
        """Enter an auto-detected handoff for a nested subgraph invocation.

        Creates a HandoffContext (HANDOFF span in parent trace + state swap)
        and stores it keyed by the subgraph root's run_id for restoration
        in on_chain_end / on_chain_error.

        Args:
            run_id: The subgraph root's run_id (used as key for exit).
            subagent_name: Display name of the target agent.
            tool_run_id: The triggering tool's run_id (for correct parenting
                under parallel tools).
        """
        ctx = HandoffContext(
            handler=self,
            target_agent=subagent_name,
            parent_tool_run_id=str(tool_run_id) if tool_run_id else None,
        )
        ctx._enter()
        self._auto_handoff_contexts[str(run_id)] = ctx
        # Store child state in handler-level dicts for cross-context access.
        # ContextVar isolation doesn't propagate across async tasks, so
        # subsequent callbacks use _ensure_auto_handoff_context() to
        # restore the right state.
        child_state = _active_handoff_state.get()
        assert child_state is not None
        run_key = str(run_id)
        self._auto_handoff_child_states[run_key] = child_state
        self._auto_handoff_run_map[run_key] = run_key
        assert ctx.handoff_span is not None
        logger.info(
            f"Auto-handoff entered: subagent={subagent_name}, run_id={run_id}, "
            f"new_trace_id={ctx.trace_id}, handoff_span_id={ctx.handoff_span.span_id}"
        )

    def _auto_handoff_exit(
        self, run_id: UUID, error: BaseException | None = None
    ) -> None:
        """Exit an auto-detected handoff, restoring parent state.

        Called when the subgraph root chain ends (on_chain_end or on_chain_error).
        """
        ctx = self._auto_handoff_contexts.pop(str(run_id), None)
        if ctx is None:
            return
        # Clean up handler-level auto-handoff mappings
        root_key = str(run_id)
        self._auto_handoff_child_states.pop(root_key, None)
        to_remove = [k for k, v in self._auto_handoff_run_map.items() if v == root_key]
        for k in to_remove:
            del self._auto_handoff_run_map[k]

        ctx._exit(error)
        s = _StateAccessor(self)
        logger.info(
            f"Auto-handoff exited: subagent={ctx._target_agent}, run_id={run_id}, "
            f"restored_trace_id={s.current_trace_id}"
        )

    def _ensure_auto_handoff_context(
        self, run_id: UUID, parent_run_id: UUID | None
    ) -> None:
        """Ensure the ContextVar points to the correct auto-handoff child state.

        LangChain fires callbacks in executor threads whose Contexts are
        copies of the event-loop Context. ContextVar changes made during
        ``_auto_handoff_enter`` (in the ``on_chain_start`` callback) are
        therefore invisible to subsequent callbacks for the same chain.

        This method uses the handler-level ``_auto_handoff_run_map`` to
        detect whether a callback belongs to an auto-handoff child graph
        and, if so, sets the ContextVar so that ``_StateAccessor`` reads
        the correct isolated state.
        """
        run_key = str(run_id)
        parent_key = str(parent_run_id) if parent_run_id else None

        # Already mapped?
        root_key = self._auto_handoff_run_map.get(run_key)
        if root_key is None and parent_key:
            root_key = self._auto_handoff_run_map.get(parent_key)
            if root_key is not None:
                # Register this run_id under the same auto-handoff root
                self._auto_handoff_run_map[run_key] = root_key

        if root_key is not None:
            child_state = self._auto_handoff_child_states.get(root_key)
            if child_state is not None:
                current = _active_handoff_state.get()
                if current is not child_state:
                    _active_handoff_state.set(child_state)

    def _serialize_for_json(self, obj: Any) -> Any:
        """Recursively serialize an object to ensure JSON compatibility.

        Args:
            obj: Object to serialize

        Returns:
            JSON-compatible representation of the object
        """
        if obj is None:
            return None

        if isinstance(obj, (str, int, float, bool)):
            return obj

        if isinstance(obj, UUID):
            return str(obj)

        if isinstance(obj, BaseMessage):
            result = {
                "role": obj.type if hasattr(obj, "type") else "unknown",
                "content": obj.content if hasattr(obj, "content") else str(obj),
            }
            if hasattr(obj, "additional_kwargs") and obj.additional_kwargs:
                result["additional_kwargs"] = self._serialize_for_json(obj.additional_kwargs)
            if hasattr(obj, "tool_calls") and obj.tool_calls:
                result["tool_calls"] = self._serialize_for_json(obj.tool_calls)
            if hasattr(obj, "tool_call_id") and obj.tool_call_id:
                result["tool_call_id"] = obj.tool_call_id
            return result

        if isinstance(obj, Document):
            return {
                "page_content": obj.page_content,
                "metadata": self._serialize_for_json(obj.metadata) if hasattr(obj, "metadata") else {},
            }

        if isinstance(obj, dict):
            return {key: self._serialize_for_json(value) for key, value in obj.items()}

        if isinstance(obj, (list, tuple)):
            return [self._serialize_for_json(item) for item in obj]

        if isinstance(obj, set):
            return [self._serialize_for_json(item) for item in obj]

        # Handle Pydantic models
        if hasattr(obj, "model_dump"):
            try:
                return self._serialize_for_json(obj.model_dump())
            except Exception:
                pass
        elif hasattr(obj, "dict"):
            try:
                return self._serialize_for_json(obj.dict())
            except Exception:
                pass

        try:
            return str(obj)
        except Exception:
            return "<unserializable>"

    def _extract_name(self, serialized: dict[str, Any], **kwargs: Any) -> str:
        """Extract component name from serialized data.

        Args:
            serialized: Component's serialized definition
            **kwargs: Additional arguments

        Returns:
            Component name string
        """
        try:
            if "name" in kwargs and kwargs["name"] is not None:
                return str(kwargs["name"])
            if "name" in serialized:
                return serialized["name"]
            if "id" in serialized and isinstance(serialized["id"], list):
                return serialized["id"][-1] if serialized["id"] else "Unknown"
            return "Unknown"
        except Exception as e:
            logger.debug(f"Error extracting component name: {e}")
            return "Unknown"

    def _has_langgraph_metadata(
        self,
        metadata: dict[str, Any] | None = None,
        tags: list[str] | None = None,
    ) -> bool:
        """Check if this span has LangGraph-specific metadata.

        Args:
            metadata: Callback metadata
            tags: Callback tags

        Returns:
            True if this span has LangGraph metadata
        """
        try:
            if metadata:
                for key in metadata.keys():
                    if key.startswith("langgraph_"):
                        return True

            if tags:
                for tag in tags:
                    if isinstance(tag, str) and tag.startswith("graph:step:"):
                        return True

            return False
        except Exception:
            return False

    def _is_langgraph_interrupt(self, error: BaseException) -> bool:
        """Check if an exception is a LangGraph Interrupt.

        LangGraph's interrupt() function raises an Interrupt exception for HITL flows.
        When propagated through callbacks, it may be wrapped in a GraphInterrupt.
        We detect this by class name to avoid a hard dependency on langgraph.

        Args:
            error: The exception to check

        Returns:
            True if this is a LangGraph Interrupt, False otherwise
        """
        error_type = type(error).__name__
        # Direct Interrupt exception
        if error_type == "Interrupt" and hasattr(error, "value"):
            return True
        # GraphInterrupt wrapper (contains interrupts list)
        if error_type == "GraphInterrupt":
            return True
        return False

    def _extract_interrupt_info(self, error: BaseException) -> dict[str, Any]:
        """Extract information from a LangGraph Interrupt or GraphInterrupt.

        Tries to cleanly extract the interrupt value (prompt string) and id
        (checkpoint ID) as separate fields. Falls back to str(error) if
        extraction fails.

        Args:
            error: The Interrupt or GraphInterrupt exception

        Returns:
            Dict with 'value' (the prompt) and 'id' (interrupt checkpoint ID)
        """
        try:
            interrupt_obj = self._find_interrupt_object(error)
            if interrupt_obj is not None:
                value = getattr(interrupt_obj, "value", None)
                interrupt_id = getattr(interrupt_obj, "id", None)
                # Ensure value is a clean scalar (not another Interrupt-like object)
                if value is not None and hasattr(value, "value"):
                    value = getattr(value, "value", str(value))
                return {
                    "value": value if value is not None else str(error),
                    "id": str(interrupt_id) if interrupt_id is not None else None,
                }
        except Exception:
            logger.debug(f"Failed to extract interrupt info, falling back to str: {error}")

        return {
            "value": str(error),
            "id": None,
        }

    def _find_interrupt_object(self, error: BaseException) -> Any:
        """Find the first Interrupt object from an error.

        Handles both direct Interrupt exceptions and GraphInterrupt wrappers.
        LangGraph stores interrupts in different ways depending on version:
        - Named attribute: error.interrupts
        - Exception args: error.args[0] (tuple of Interrupt objects)

        Returns:
            The Interrupt object, or None if not found
        """
        error_type = type(error).__name__

        # Direct Interrupt exception
        if error_type == "Interrupt" and hasattr(error, "value"):
            return error

        # GraphInterrupt wrapper - extract first interrupt
        if error_type == "GraphInterrupt":
            # Try named attribute first
            interrupts = getattr(error, "interrupts", None)
            if interrupts:
                first = interrupts[0]
                if hasattr(first, "value"):
                    return first

            # Fall back to error.args[0] (tuple of Interrupt objects)
            if error.args:
                first_arg = error.args[0]
                # args[0] may be a tuple/list of Interrupt objects
                if isinstance(first_arg, (tuple, list)) and first_arg:
                    candidate = first_arg[0]
                    if hasattr(candidate, "value"):
                        return candidate
                # args[0] may be the Interrupt object directly
                if hasattr(first_arg, "value"):
                    return first_arg

        return None

    def _extract_resume_value(self, inputs: dict[str, Any] | Any) -> str | None:
        """Extract the clean resume value from LangGraph Command inputs.

        LangGraph resumes pass inputs containing a Command(resume=...) object.
        This method extracts the resume value from that Command, falling back
        to str() of the serialized inputs.

        Args:
            inputs: The inputs dict from on_chain_start

        Returns:
            The resume value as a string, or None if no inputs
        """
        if not inputs:
            return None

        try:
            command = self._find_command_object(inputs)
            if command is not None:
                resume = getattr(command, "resume", None)
                if resume is not None:
                    serialized = self._serialize_for_json(resume)
                    if isinstance(serialized, str):
                        return serialized
                    return json.dumps(serialized, default=str)
        except Exception:
            logger.debug("Failed to extract resume value from Command, falling back to str")

        return str(self._serialize_for_json(inputs))

    def _find_command_object(self, obj: Any) -> Any:
        """Find a LangGraph Command object in inputs.

        Args:
            obj: The object to search (dict, list, or direct Command)

        Returns:
            The Command object, or None if not found
        """
        if type(obj).__name__ == "Command" and hasattr(obj, "resume"):
            return obj

        if isinstance(obj, dict):
            for value in obj.values():
                found = self._find_command_object(value)
                if found is not None:
                    return found

        if isinstance(obj, (list, tuple)):
            for item in obj:
                found = self._find_command_object(item)
                if found is not None:
                    return found

        return None

    def _is_cancelled_error(self, error: BaseException) -> bool:
        """Check if an exception is a CancelledError (user cancelled the stream).

        This occurs when the user cancels an in-progress agent execution via
        the UI or API. Unlike interrupts (HITL), this is not a deliberate pause
        but a user-initiated abort.

        Args:
            error: The exception to check

        Returns:
            True if this is a CancelledError, False otherwise
        """
        import asyncio
        # Check for asyncio.CancelledError or any CancelledError variant
        if isinstance(error, asyncio.CancelledError):
            return True
        # Also check by name for wrapped exceptions
        return type(error).__name__ == "CancelledError"

    def _extract_provider_name(self, serialized: dict[str, Any], kwargs: dict[str, Any]) -> str | None:
        """Extract the LLM provider name from serialized data or invocation params."""
        if "id" in serialized and isinstance(serialized["id"], list):
            for part in serialized["id"]:
                part_lower = str(part).lower()
                for provider in ("openai", "anthropic", "google", "gemini", "bedrock", "cohere", "mistral"):
                    if provider in part_lower:
                        return "google" if provider == "gemini" else provider

        invocation_params = kwargs.get("invocation_params", {})
        model_type = invocation_params.get("_type", "")
        if model_type:
            type_lower = model_type.lower()
            for provider in ("openai", "anthropic", "google", "bedrock", "cohere", "mistral"):
                if provider in type_lower:
                    return provider

        return None

    def _set_model_parameters(self, span: Span, kwargs: dict[str, Any]) -> None:
        """Extract and set model parameters as standard attributes.

        Args:
            span: The span to set attributes on
            kwargs: The kwargs dict from serialized model data
        """
        # Map LangChain parameter names to OTEL standard attributes
        param_mapping = {
            "temperature": self.ATTR_REQUEST_TEMPERATURE,
            "top_p": self.ATTR_REQUEST_TOP_P,
            "top_k": self.ATTR_REQUEST_TOP_K,
            "max_tokens": self.ATTR_REQUEST_MAX_TOKENS,
            "max_output_tokens": self.ATTR_REQUEST_MAX_TOKENS,  # Alternative name
            "frequency_penalty": self.ATTR_REQUEST_FREQUENCY_PENALTY,
            "presence_penalty": self.ATTR_REQUEST_PRESENCE_PENALTY,
        }

        for param_name, attr_name in param_mapping.items():
            if param_name in kwargs and kwargs[param_name] is not None:
                span.set_attribute(attr_name, kwargs[param_name])

        # Handle list/array parameters separately (need JSON serialization)
        for param_name in ("stop", "stop_sequences"):
            value = kwargs.get(param_name)
            if value is not None:
                span.set_attribute(
                    self.ATTR_REQUEST_STOP_SEQUENCES,
                    json.dumps(value) if isinstance(value, (list, tuple)) else value,
                )
                break

        seed = kwargs.get("seed")
        if seed is not None:
            span.set_attribute(self.ATTR_REQUEST_SEED, seed)

    def _start_span(
        self,
        run_id: UUID,
        parent_run_id: UUID | None,
        name: str,
        span_kind: SpanKind,
        span_type: SpanType,
    ) -> None:
        """Start a new SDK span for a LangChain operation.

        Args:
            run_id: Unique run identifier
            parent_run_id: Parent run identifier (for nesting)
            name: Span name
            span_kind: Span kind
            span_type: Beacon span type (SpanType enum)
        """
        s = _StateAccessor(self)

        # Determine parent from LangChain's parent_run_id
        parent_id = None
        trace_id = None
        if parent_run_id:
            parent_key = str(parent_run_id)
            if parent_key in s.runs:
                parent_span = s.runs[parent_key]["span"]
                parent_id = parent_span.span_id
                trace_id = parent_span.trace_id
                logger.debug(
                    f"_start_span found parent: name={name}, run_id={run_id}, "
                    f"parent_run_id={parent_run_id}, parent_span_id={parent_id}"
                )

        # If no parent trace, use current trace_id (for LangGraph resume or continuing trace)
        if trace_id is None:
            trace_id = s.current_trace_id

        # If resume mode and no parent, attach under the root span
        # This flattens resume spans under the original root span
        if parent_id is None and s.is_resume and s.root_span_id:
            parent_id = s.root_span_id
            logger.debug(
                f"_start_span using root_span_id (resume mode): name={name}, "
                f"parent_id={parent_id}, run_id={run_id}"
            )

        # Log orphaned spans (no parent found)
        if parent_id is None and parent_run_id is not None:
            logger.debug(
                f"_start_span orphaned: name={name}, run_id={run_id}, "
                f"parent_run_id={parent_run_id} not in runs"
            )

        # Save previous span for context restoration
        previous_span = get_current_span()

        # Get timestamp override from client (for historical data seeding)
        timestamp_override = self.client._timestamp_override

        # Track whether this is the first span (root) before we add to _runs
        is_first_span = not s.runs
        if is_first_span:
            s.interrupt_detected = False
            s.interrupt_info = {}

        # Use pre-generated root_span_id for first span if set (and not in resume mode)
        span_id = None
        if s.root_span_id and not s.is_resume and is_first_span:
            # First span with pre-generated ID - use it
            span_id = s.root_span_id
            logger.debug(f"Using pre-generated root_span_id={span_id} for first span")

        # Override span type to AGENT for the handoff child's root span only.
        # Cannot rely on is_first_span (= not s.runs) because s.runs empties
        # as completed spans are deleted, making later top-level nodes appear
        # "first" when the resume root was skipped and never added to runs.
        is_handoff_root = False
        hs = _active_handoff_state.get()
        if hs and s.parent_handoff_trace_id and not hs.root_span_created:
            span_type = SpanType.AGENT
            hs.root_span_created = True
            is_handoff_root = True

        # Create SDK Span
        span = Span(
            name=name,
            trace_id=trace_id,  # None = generate new trace
            span_id=span_id,  # None = generate new, or use pre-generated
            parent_id=parent_id,
            kind=span_kind,
            span_type=span_type,
            session_id=self._session_id,
            user_email=self._user_email,
            user_name=self._user_name,
            user_id=self._user_id,
            record_stacktrace=self._record_stacktrace,
        )
        span.start(timestamp_override)

        # Set OTEL conversation ID alongside beacon session_id
        if self._session_id:
            span.set_attribute(self.ATTR_CONVERSATION_ID, self._session_id)

        # Set common attributes
        span.set_attribute(self.ATTR_COMPONENT_TYPE, name)

        # Set deployment environment
        if self._environment:
            span.set_attribute(self.ATTR_DEPLOYMENT_ENVIRONMENT, self._environment)

        # Set agent attributes
        span.set_attribute(self.ATTR_AGENT_FRAMEWORK, "langgraph")
        if s.agent_name:
            span.set_attribute(self.ATTR_AGENT_NAME, s.agent_name)

        # Set custom metadata
        if self._metadata:
            for key, value in self._metadata.items():
                attr_key = f"{self.ATTR_BEACON_METADATA_PREFIX}.{key}"
                span.set_attribute(attr_key, value)

        # Set as current span using SDK's ContextVar (async-safe)
        set_current_span(span)

        # Store span and context info
        s.runs[str(run_id)] = {
            "span": span,
            "parent_run_id": str(parent_run_id) if parent_run_id else None,
            "previous_span": previous_span,
            "start_perf_time": time.perf_counter(),
            "timestamp_override": timestamp_override,
        }

        # Update current trace_id (for trace continuation support)
        if s.current_trace_id is None:
            s.current_trace_id = span.trace_id

        # Capture root span ID for resume (LangGraph reruns from top)
        if s.root_span_id is None:
            s.root_span_id = span.span_id

        # Add back-link and handoff attributes to child root span
        if is_handoff_root and s.parent_handoff_span_id:
            span.add_link(
                trace_id=s.parent_handoff_trace_id,
                span_id=s.parent_handoff_span_id,
                attributes={"link.type": "parent_handoff"},
            )
            span.set_attribute("handoff.source_trace_id", s.parent_handoff_trace_id)
            span.set_attribute("handoff.source_span_id", s.parent_handoff_span_id)
            if s.parent_handoff_agent_name:
                span.set_attribute("handoff.source_agent_name", s.parent_handoff_agent_name)

        # Eagerly export span metadata so it appears in the UI immediately
        try:
            self.client.export_eager_span(span)
        except Exception:
            pass  # Eager export failure must never break span execution

        logger.debug(
            f"Started LangChain span: name={name}, type={span_type.value}, "
            f"trace_id={span.trace_id}, span_id={span.span_id}"
        )

    def _end_span(self, run_id: UUID, error: BaseException | None = None) -> None:
        """End an SDK span.

        Args:
            run_id: Run identifier of the span to end
            error: Optional exception if the operation failed
        """
        s = _StateAccessor(self)
        run_id_str = str(run_id)
        if run_id_str not in s.runs:
            return

        run_data = s.runs[run_id_str]
        span = run_data["span"]

        # Set status and record error if present
        if error:
            if self._is_langgraph_interrupt(error):
                # Handle LangGraph interrupt - NOT an error, this is normal HITL flow
                # All spans in the chain get marked as interrupted (they all were)
                # Resume will attach under root since LangGraph reruns from top
                interrupt_info = self._extract_interrupt_info(error)
                logger.info(
                    f"LangGraph interrupt detected: value={interrupt_info.get('value')}, "
                    f"id={interrupt_info.get('id')}, span={span.name}"
                )
                span.set_attribute(self.ATTR_LANGGRAPH_INTERRUPT, True)
                if interrupt_info.get("id"):
                    span.set_attribute(self.ATTR_LANGGRAPH_INTERRUPT_ID, interrupt_info["id"])
                if interrupt_info.get("value"):
                    span.set_attribute(self.ATTR_LANGGRAPH_INTERRUPT_VALUE, str(interrupt_info["value"]))
                span.set_status(StatusCode.OK, f"Interrupted: {interrupt_info.get('value', 'awaiting input')}")
                # Track interrupt at handler level for root span propagation
                s.interrupt_detected = True
                s.interrupt_info = interrupt_info
                # Persist interrupt ID at class level so resume handlers can pick it up
                if interrupt_info.get("id") and s.current_trace_id:
                    BeaconLangGraphHandler._trace_interrupt_ids[s.current_trace_id] = interrupt_info["id"]
            elif self._is_cancelled_error(error):
                # Handle user cancellation - NOT an error, user aborted the stream
                # All spans in the chain get marked as cancelled (they all were)
                # Resume will attach under root since LangGraph reruns from top
                logger.info(f"User cancellation detected: span={span.name} (marked as cancelled)")
                span.set_attribute(self.ATTR_LANGGRAPH_CANCELLED, True)
                span.set_status(StatusCode.OK, "Cancelled by user")

                # Attach partial streaming output from this span or any child LLM span
                try:
                    partial_text, partial_thinking = self._collect_partial_streaming_output()
                    if partial_text or partial_thinking:
                        content_blocks: list[dict[str, str]] = []
                        if partial_thinking:
                            content_blocks.append({"type": "thinking", "thinking": partial_thinking})
                        if partial_text:
                            content_blocks.append({"type": "text", "text": partial_text})
                        message = {"role": "assistant", "content": content_blocks}
                        output = {"messages": [message]}
                        span.set_attribute(
                            self.ATTR_LLM_OUTPUT_MESSAGES,
                            json.dumps(output, default=str),
                        )
                except Exception as e:
                    logger.debug(f"Failed to collect partial streaming output: {e}")
            else:
                span.set_attribute("error.type", type(error).__qualname__)
                span.record_exception(error)
                span.set_status(StatusCode.ERROR, str(error))
        else:
            if span.status_code == StatusCode.UNSET:
                span.set_status(StatusCode.OK)

        # Calculate end_time if using timestamp override
        end_time = None
        timestamp_override = run_data.get("timestamp_override")
        if timestamp_override is not None:
            start_perf_time = run_data.get("start_perf_time")
            if start_perf_time is not None:
                duration_seconds = time.perf_counter() - start_perf_time
                end_time = timestamp_override + timedelta(seconds=duration_seconds)

        span.end(end_time=end_time)

        # Export span via SDK's mechanism
        export_result = self.client.export_span(span)
        logger.debug(
            f"Ended LangChain span: name={span.name}, "
            f"trace_id={span.trace_id}, span_id={span.span_id}, "
            f"export_success={export_result}"
        )

        # Restore previous context
        previous_span = run_data.get("previous_span")
        if previous_span is not None:
            set_current_span(previous_span)
        elif run_data.get("parent_run_id") is None:
            # Root span ended - clear context entirely
            clear_context()
            # Only mark invocation completed for the handler's base state,
            # not for auto-handoff child traces (whose root spans also have
            # parent_run_id=None but run in isolated ContextVar state).
            if not _active_handoff_state.get():
                self._invocation_completed = True

        del s.runs[run_id_str]

        # When root span ends, flush any orphaned child spans
        # (e.g., LLM span eagerly exported but never got on_llm_end due to CancelledError)
        if run_data.get("parent_run_id") is None and s.runs:
            self._flush_orphaned_spans()

    def _set_token_usage(self, span: Any, usage: dict) -> None:
        """Set token usage attributes on a span from a usage metadata dict.

        Handles both basic counts (input/output/total) and detailed
        breakdowns (reasoning tokens, cache tokens) when available.
        """
        input_tokens = usage.get("input_tokens", 0) or usage.get("prompt_tokens", 0) or 0
        output_tokens = usage.get("output_tokens", 0) or usage.get("completion_tokens", 0) or 0

        span.set_attribute(self.ATTR_LLM_INPUT_TOKENS, input_tokens)
        span.set_attribute(self.ATTR_LLM_OUTPUT_TOKENS, output_tokens)

        # Detailed breakdowns (Anthropic thinking, prompt caching)
        output_details = usage.get("output_token_details") or {}
        reasoning_tokens = output_details.get("reasoning", 0) or 0
        if reasoning_tokens:
            span.set_attribute(self.ATTR_LLM_REASONING_TOKENS, reasoning_tokens)

        input_details = usage.get("input_token_details") or {}
        cache_read = input_details.get("cache_read", 0) or 0
        cache_creation = input_details.get("cache_creation", 0) or 0
        if cache_read:
            span.set_attribute(self.ATTR_LLM_CACHE_READ_TOKENS, cache_read)
        if cache_creation:
            span.set_attribute(self.ATTR_LLM_CACHE_CREATION_TOKENS, cache_creation)

    def _flush_orphaned_spans(self) -> None:
        """End and export all remaining spans as cancelled.

        Called when the root span ends but child spans are still active,
        which happens when CancelledError interrupts astream_events
        before LangChain callbacks fire for in-progress LLM calls.

        Preserves any data already accumulated during streaming:
        partial output, token usage, TTFT, and model name.
        """
        s = _StateAccessor(self)
        for run_id_str in list(s.runs):
            run_data = s.runs[run_id_str]
            span = run_data["span"]
            span.set_attribute(self.ATTR_LANGGRAPH_CANCELLED, True)
            span.set_status(StatusCode.OK, "Cancelled by user")

            # Attach partial streaming output
            text_parts = run_data.get("streaming_content")
            thinking_parts = run_data.get("streaming_thinking")
            if text_parts or thinking_parts:
                content_blocks: list[dict[str, str]] = []
                if thinking_parts:
                    content_blocks.append({"type": "thinking", "thinking": "".join(thinking_parts)})
                if text_parts:
                    content_blocks.append({"type": "text", "text": "".join(text_parts)})
                span.set_attribute(
                    self.ATTR_LLM_OUTPUT_MESSAGES,
                    json.dumps({"messages": [{"role": "assistant", "content": content_blocks}]}, default=str),
                )

            # Attach token usage captured from streaming chunks
            streaming_usage = run_data.get("streaming_usage")
            if streaming_usage:
                self._set_token_usage(span, streaming_usage)

            # Attach TTFT if first token was received
            first_token_time = s.first_token_times.pop(run_id_str, None)
            if first_token_time is not None:
                start_perf_time = run_data.get("start_perf_time")
                if start_perf_time is not None:
                    ttft_ms = (first_token_time - start_perf_time) * 1000
                    span.set_attribute(self.ATTR_TTFT_MS, ttft_ms)
                    span.set_attribute(self.ATTR_STREAMING, True)

            # Attach model name if cached
            model_name = run_data.get("cached_model_name")
            if model_name:
                span.set_attribute(self.ATTR_LLM_MODEL, model_name)
                span.set_attribute(self.ATTR_RESPONSE_MODEL, model_name)

            span.end()
            self.client.export_span(span)
            logger.debug(f"Flushed orphaned span: name={span.name}, span_id={span.span_id}")
        s.runs.clear()

    def _collect_partial_streaming_output(self) -> tuple[str | None, str | None]:
        """Collect partial streaming output from any active LLM span.

        When a cancellation happens mid-stream, the LLM span may have accumulated
        tokens via on_llm_new_token but never received on_llm_end. This method
        finds that partial content so it can be attached to the cancelled span.

        Returns:
            Tuple of (text content, thinking content). Either or both may be None.
        """
        # Search all active spans for streaming content (the LLM span
        # that was mid-stream could be a direct child or nested deeper)
        s = _StateAccessor(self)
        for _, data in s.runs.items():
            text = data.get("streaming_content")
            thinking = data.get("streaming_thinking")
            if text or thinking:
                return (
                    "".join(text) if text else None,
                    "".join(thinking) if thinking else None,
                )

        return None, None

    def _update_root_span_with_resume_output(self, outputs: dict[str, Any]) -> None:
        """Re-export the root span with the final output after resume completion.

        When a LangGraph agent is interrupted/cancelled and then resumed, the
        original root span was already exported without output. This method
        creates an update span with the same trace_id + span_id so the backend
        can merge the final output into the original root span.
        """
        s = _StateAccessor(self)
        if not s.root_span_id or not s.current_trace_id:
            return

        root_span = Span(
            name=s.agent_name or "agent",
            trace_id=s.current_trace_id,
            span_id=s.root_span_id,
            span_type=SpanType.AGENT,
            session_id=self._session_id,
            user_email=self._user_email,
            user_name=self._user_name,
            user_id=self._user_id,
        )

        # Set the final output
        serialized = self._serialize_for_json(outputs)
        root_span.set_attribute(
            self.ATTR_LLM_OUTPUT_MESSAGES,
            json.dumps(serialized, default=str),
        )

        # Set agent attributes
        root_span.set_attribute(self.ATTR_AGENT_FRAMEWORK, "langgraph")
        if s.agent_name:
            root_span.set_attribute(self.ATTR_AGENT_NAME, s.agent_name)

        # Override cancelled/interrupted status (unless a new interrupt was detected during resume)
        root_span.set_attribute(self.ATTR_LANGGRAPH_CANCELLED, False)
        root_span.set_attribute("langgraph.resumed", True)
        if s.interrupt_detected:
            root_span.set_attribute(self.ATTR_LANGGRAPH_INTERRUPT, True)
            if s.interrupt_info.get("id"):
                root_span.set_attribute(self.ATTR_LANGGRAPH_INTERRUPT_ID, s.interrupt_info["id"])
            if s.interrupt_info.get("value"):
                root_span.set_attribute(
                    self.ATTR_LANGGRAPH_INTERRUPT_VALUE,
                    str(s.interrupt_info["value"]),
                )
            root_span.set_status(
                StatusCode.OK,
                f"Interrupted: {s.interrupt_info.get('value', 'awaiting input')}",
            )
        else:
            root_span.set_attribute(self.ATTR_LANGGRAPH_INTERRUPT, False)
            root_span.set_status(StatusCode.OK)
            # Clean up class-level interrupt ID cache when trace completes without new interrupt
            if s.current_trace_id:
                self._trace_interrupt_ids.pop(s.current_trace_id, None)

        root_span.start()
        root_span.end()
        self.client.export_span(root_span)

        logger.debug(
            f"Resume root span exported: trace_id={s.current_trace_id}, "
            f"span_id={s.root_span_id}"
        )

    # Chain callbacks

    def on_chain_start(
        self,
        serialized: dict[str, Any],
        inputs: dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle chain start event."""
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            s = _StateAccessor(self)

            # Inside a manual handoff context, the child graph's root chain
            # arrives with a parent_run_id inherited from the parent context
            # (e.g. the orchestrator's tool run_id). Clear it so the child's
            # root is correctly detected for resume skipping and root span
            # creation. Only clear when the parent is truly external — not
            # in the child's runs or skipped_resume_roots.
            if (
                s.in_handoff
                and parent_run_id is not None
                and str(parent_run_id) not in s.runs
                and str(parent_run_id) not in s.skipped_resume_roots
            ):
                parent_run_id = None

            # Reset per-invocation state when reusing handler across multiple calls.
            # Conditions: root span (no parent), not a resume, not inside a handoff,
            # and a previous invocation has completed on this handler.
            if (
                parent_run_id is None
                and not s.is_resume
                and not s.in_handoff
                and self._invocation_completed
            ):
                logger.debug(
                    "Resetting handler state for new invocation "
                    f"(previous trace_id={self._current_trace_id})"
                )
                self._reset_invocation_state()
                s = _StateAccessor(self)

            name = self._extract_name(serialized, **kwargs)

            has_lg_metadata = self._has_langgraph_metadata(metadata=metadata, tags=tags)

            # Check for LangGraph metadata
            if parent_run_id and has_lg_metadata:
                s.langgraph_parent_ids.add(str(parent_run_id))

            # Extract thread_id from metadata for attribute setting
            thread_id = None
            if metadata:
                thread_id = metadata.get("langgraph_thread_id") or metadata.get("thread_id")

            # Detect subagent roots: a compiled graph invoked from within a
            # tool has parent_run_id pointing to the active tool.
            # Chains *inside* the subgraph point to graph-internal parents.
            is_subagent_root = (
                has_lg_metadata
                and parent_run_id is not None
                and str(parent_run_id) in s.active_tool_run_ids
            )

            if is_subagent_root and self._autodetect_handoffs:
                subagent_name = (metadata or {}).get("lc_agent_name") or name
                logger.info(
                    f"Auto-handoff detected: name={subagent_name}, run_id={run_id}, "
                    f"parent_tool_run_id={parent_run_id}"
                )
                self._auto_handoff_enter(run_id, subagent_name, tool_run_id=parent_run_id)
                # Clear parent so subgraph root becomes trace root in child trace
                parent_run_id = None
                # Re-acquire state accessor after auto-handoff changed ContextVar
                s = _StateAccessor(self)

            # If no parent but there's an active tool, re-parent under the tool
            if parent_run_id is None and s.active_tool_run_ids:
                parent_run_id = UUID(s.active_tool_run_ids[-1])
                logger.info(
                    f"Using active tool as parent for nested graph: name={name}, "
                    f"parent_tool_run_id={parent_run_id}"
                )

            # Skip the resume graph span entirely - reparent children under original root
            # Only skip the actual resume root (no parent + we're in resume mode)
            should_skip_as_resume_root = parent_run_id is None and s.is_resume

            if should_skip_as_resume_root:
                # This is a resume graph span - don't create a span for it
                # Capture the resume value from inputs (e.g. Command(resume=...))
                resume_value = self._extract_resume_value(inputs)
                # Look up interrupt ID: constructor param > instance state > class-level cache
                resolved_interrupt_id = (
                    self._interrupt_id
                    or s.interrupt_info.get("id")
                    or (self._trace_interrupt_ids.get(s.current_trace_id) if s.current_trace_id else None)
                )
                s.skipped_resume_roots[str(run_id)] = {
                    "parent_id": s.root_span_id,
                    "resume_value": resume_value,
                    "interrupt_id": resolved_interrupt_id,
                }
                return

            # If parent was a skipped resume root, reparent under the original root
            is_first_resume_child = False
            resume_value = None
            if parent_run_id and str(parent_run_id) in s.skipped_resume_roots:
                skipped_info = s.skipped_resume_roots[str(parent_run_id)]
                # Only mark the first child as resume
                if not skipped_info.get("first_child_marked"):
                    resume_value = skipped_info.get("resume_value")
                    is_first_resume_child = True
                    skipped_info["first_child_marked"] = True
                # Override parent_run_id to None so _start_span uses root_span_id (resume mode)
                parent_run_id = None

            # Subagent roots get AGENT type even without full handoff
            chain_type = SpanType.AGENT if (is_subagent_root and not self._autodetect_handoffs) else SpanType.CHAIN
            self._start_span(run_id, parent_run_id, name, SpanKind.INTERNAL, chain_type)

            # Track potential conditional routing functions for reclassification in on_chain_end.
            # Routing functions run as internal steps within a node's RunnableSeq (seq:step tags)
            # rather than as top-level graph nodes (graph:step tags).
            has_seq_step = tags and any(
                isinstance(t, str) and t.startswith("seq:step:") for t in tags
            )
            has_graph_step = tags and any(
                isinstance(t, str) and t.startswith("graph:step:") for t in tags
            )
            if (
                has_seq_step
                and not has_graph_step
                and has_lg_metadata
                and name not in self._LANGGRAPH_INTERNAL_NAMES
            ):
                s.potential_conditional_run_ids.add(str(run_id))

            if str(run_id) in s.runs:
                span = s.runs[str(run_id)]["span"]
                serialized_inputs = self._serialize_for_json(inputs)
                span.set_attribute(self.ATTR_LLM_INPUT_MESSAGES, json.dumps(serialized_inputs, default=str))

                # Mark first resume span with resume attributes
                if is_first_resume_child:
                    span.set_attribute(self.ATTR_LANGGRAPH_RESUME, True)
                    if resume_value:
                        span.set_attribute(self.ATTR_LANGGRAPH_RESUME_VALUE, resume_value)
                    interrupt_id = skipped_info.get("interrupt_id")
                    if interrupt_id:
                        span.set_attribute(self.ATTR_LANGGRAPH_INTERRUPT_ID, interrupt_id)

                # Set operation name for agent root spans
                if parent_run_id is None and not is_first_resume_child:
                    span.set_attribute(self.ATTR_OPERATION_NAME, "invoke_agent")

                # Set metadata as attributes (skip beacon checkpoint keys —
                # they duplicate data already present on the span itself).
                if metadata:
                    for key, value in metadata.items():
                        if key in _BEACON_CHECKPOINT_KEYS:
                            continue
                        if isinstance(value, (str, int, float, bool)):
                            span.set_attribute(f"langchain.metadata.{key}", value)
                        else:
                            span.set_attribute(f"langchain.metadata.{key}", json.dumps(value, default=str))

                if tags:
                    span.set_attribute("langchain.tags", json.dumps(tags))

                # Try to extract graph name from LangGraph metadata if agent_name not set
                if not s.agent_name and metadata:
                    graph_name = metadata.get("langgraph_checkpoint_ns") or metadata.get("name")
                    if graph_name:
                        s.agent_name = graph_name

                # Set thread_id attribute on root spans
                if thread_id and parent_run_id is None and not is_first_resume_child:
                    span.set_attribute(self.ATTR_LANGGRAPH_THREAD_ID, thread_id)

        except Exception as e:
            logger.error(f"Error in on_chain_start: {e}")

    def on_chain_end(
        self,
        outputs: dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle chain end event."""
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            s = _StateAccessor(self)

            # If this was a skipped resume root, re-export root span with output
            if str(run_id) in s.skipped_resume_roots:
                self._update_root_span_with_resume_output(outputs)
                del s.skipped_resume_roots[str(run_id)]
                return

            if str(run_id) in s.runs:
                span = s.runs[str(run_id)]["span"]
                serialized_outputs = self._serialize_for_json(outputs)
                if serialized_outputs is not None:
                    span.set_attribute(self.ATTR_LLM_OUTPUT_MESSAGES, json.dumps(serialized_outputs, default=str))

                # Reclassify conditional routing functions
                if str(run_id) in s.potential_conditional_run_ids:
                    s.potential_conditional_run_ids.discard(str(run_id))
                    # Routing functions return a str (destination node) or list of Send objects
                    is_routing_output = isinstance(outputs, str) or (
                        isinstance(outputs, list)
                        and all(hasattr(o, "node") for o in outputs)
                    )
                    if is_routing_output:
                        span.span_type = SpanType.CONDITIONAL
                        span.set_attribute("span.type", SpanType.CONDITIONAL.value)
                        span.set_attribute(
                            self.ATTR_LANGGRAPH_ROUTING_DECISION, str(outputs)
                        )

                # Handle LangGraph parent detection
                is_in_langgraph_parents = str(run_id) in s.langgraph_parent_ids
                if is_in_langgraph_parents:
                    if parent_run_id is None:
                        # Update both the span_type property and the span.type attribute
                        span.span_type = SpanType.AGENT
                        span.set_attribute("span.type", SpanType.AGENT.value)
                    s.langgraph_parent_ids.discard(str(run_id))

                # Propagate interrupt to root span if a child detected one
                # LangGraph may call on_chain_end (not on_chain_error) for the root
                # when an interrupt occurs, since it considers the execution complete
                if s.interrupt_detected and parent_run_id is None:
                    span.set_attribute(self.ATTR_LANGGRAPH_INTERRUPT, True)
                    if s.interrupt_info.get("id"):
                        span.set_attribute(self.ATTR_LANGGRAPH_INTERRUPT_ID, s.interrupt_info["id"])
                    if s.interrupt_info.get("value"):
                        span.set_attribute(
                            self.ATTR_LANGGRAPH_INTERRUPT_VALUE,
                            str(s.interrupt_info["value"]),
                        )
                    span.set_status(
                        StatusCode.OK,
                        f"Interrupted: {s.interrupt_info.get('value', 'awaiting input')}",
                    )

            self._end_span(run_id)

            # Auto-handoff: if this was the subgraph root, restore parent state
            if str(run_id) in self._auto_handoff_contexts:
                self._auto_handoff_exit(run_id)

        except Exception as e:
            logger.error(f"Error in on_chain_end: {e}")

    def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle chain error event."""
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            s = _StateAccessor(self)

            # Skip if this was a skipped resume root span
            if str(run_id) in s.skipped_resume_roots:
                # Resume root was cancelled — re-export root span with partial output
                # (keep cancelled=True since the resume itself was cancelled)
                if self._is_cancelled_error(error) and s.root_span_id and s.current_trace_id:
                    try:
                        partial_text, partial_thinking = self._collect_partial_streaming_output()
                        if partial_text or partial_thinking:
                            content_blocks: list[dict[str, str]] = []
                            if partial_thinking:
                                content_blocks.append({"type": "thinking", "thinking": partial_thinking})
                            if partial_text:
                                content_blocks.append({"type": "text", "text": partial_text})
                            root_span = Span(
                                name=s.agent_name or "agent",
                                trace_id=s.current_trace_id,
                                span_id=s.root_span_id,
                                span_type=SpanType.AGENT,
                                session_id=self._session_id,
                                user_email=self._user_email,
                                user_name=self._user_name,
                                user_id=self._user_id,
                            )
                            root_span.set_attribute(
                                self.ATTR_LLM_OUTPUT_MESSAGES,
                                json.dumps({"messages": [{"role": "assistant", "content": content_blocks}]}, default=str),
                            )
                            root_span.set_attribute(self.ATTR_AGENT_FRAMEWORK, "langgraph")
                            if s.agent_name:
                                root_span.set_attribute(self.ATTR_AGENT_NAME, s.agent_name)
                            root_span.set_attribute(self.ATTR_LANGGRAPH_CANCELLED, True)
                            root_span.set_attribute("langgraph.resumed", True)
                            root_span.set_status(StatusCode.OK, "Cancelled by user")
                            root_span.start()
                            root_span.end()
                            self.client.export_span(root_span)
                    except Exception as e:
                        logger.debug(f"Failed to collect partial output for cancelled resume: {e}")
                elif self._is_langgraph_interrupt(error) and s.root_span_id and s.current_trace_id:
                    # Resume root was interrupted again (multi-step HITL)
                    try:
                        interrupt_info = self._extract_interrupt_info(error)
                        root_span = Span(
                            name=s.agent_name or "agent",
                            trace_id=s.current_trace_id,
                            span_id=s.root_span_id,
                            span_type=SpanType.AGENT,
                            session_id=self._session_id,
                            user_email=self._user_email,
                            user_name=self._user_name,
                            user_id=self._user_id,
                        )
                        root_span.set_attribute(self.ATTR_AGENT_FRAMEWORK, "langgraph")
                        if s.agent_name:
                            root_span.set_attribute(self.ATTR_AGENT_NAME, s.agent_name)
                        root_span.set_attribute(self.ATTR_LANGGRAPH_INTERRUPT, True)
                        if interrupt_info.get("id"):
                            root_span.set_attribute(self.ATTR_LANGGRAPH_INTERRUPT_ID, interrupt_info["id"])
                        if interrupt_info.get("value"):
                            root_span.set_attribute(
                                self.ATTR_LANGGRAPH_INTERRUPT_VALUE,
                                str(interrupt_info["value"]),
                            )
                        root_span.set_attribute("langgraph.resumed", True)
                        root_span.set_status(
                            StatusCode.OK,
                            f"Interrupted: {interrupt_info.get('value', 'awaiting input')}",
                        )
                        root_span.start()
                        root_span.end()
                        self.client.export_span(root_span)
                    except Exception as e:
                        logger.debug(f"Failed to export interrupt status for resumed root: {e}")
                del s.skipped_resume_roots[str(run_id)]
                return

            is_interrupt = self._is_langgraph_interrupt(error)
            is_cancelled = self._is_cancelled_error(error)
            if is_interrupt:
                interrupt_info = self._extract_interrupt_info(error)

            # Handle LangGraph parent detection - set AGENT type even on error/interrupt/cancel
            # (mirrors the logic in on_chain_end)
            if str(run_id) in s.runs:
                span = s.runs[str(run_id)]["span"]
                is_in_langgraph_parents = str(run_id) in s.langgraph_parent_ids
                if is_in_langgraph_parents and parent_run_id is None:
                    span.span_type = SpanType.AGENT
                    span.set_attribute("span.type", SpanType.AGENT.value)

            # Force-end any active tool spans before ending the chain
            # This handles the case where a tool invokes a nested graph that gets cancelled
            # - the nested graph ends via on_chain_error, but on_tool_end never fires
            if is_cancelled and s.active_tool_run_ids:
                for tool_run_id in list(s.active_tool_run_ids):
                    if tool_run_id in s.runs:
                        self._end_span(UUID(tool_run_id), error)
                s.active_tool_run_ids.clear()

            s.langgraph_parent_ids.discard(str(run_id))
            self._end_span(run_id, error)

            # Auto-handoff: if this was the subgraph root, restore parent state
            if str(run_id) in self._auto_handoff_contexts:
                self._auto_handoff_exit(run_id, error)
        except Exception as e:
            logger.error(f"Error in on_chain_error: {e}")

    # LLM callbacks

    def on_llm_start(
        self,
        serialized: dict[str, Any],
        prompts: list[str],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle LLM start event (for non-chat models)."""
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            s = _StateAccessor(self)
            name = self._extract_name(serialized, **kwargs)
            self._start_span(run_id, parent_run_id, name, SpanKind.CLIENT, SpanType.GENERATION)

            # Initialize TTFT tracking (None means waiting for first token)
            s.first_token_times[str(run_id)] = None

            if str(run_id) in s.runs:
                span = s.runs[str(run_id)]["span"]
                span.set_attribute(self.ATTR_OPERATION_NAME, "chat")

                # Provider name (try to extract real provider, fallback to "langchain")
                provider = self._extract_provider_name(serialized, kwargs) or "langchain"
                span.set_attribute(self.ATTR_LLM_PROVIDER_NAME, provider)

                # Input messages (OTEL standard)
                prompt_json = json.dumps(prompts, default=str)
                span.set_attribute(self.ATTR_LLM_INPUT_MESSAGES, prompt_json)

                # Extract model name and parameters from serialized kwargs
                if "kwargs" in serialized:
                    model_kwargs = serialized["kwargs"]
                    model_name = model_kwargs.get("model_name") or model_kwargs.get("model")
                    if model_name:
                        span.set_attribute(self.ATTR_LLM_MODEL, model_name)
                        s.runs[str(run_id)]["cached_model_name"] = model_name
                    self._set_model_parameters(span, model_kwargs)

        except Exception as e:
            logger.error(f"Error in on_llm_start: {e}")

    def on_chat_model_start(
        self,
        serialized: dict[str, Any],
        messages: list[list[BaseMessage]],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle chat model start event."""
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            s = _StateAccessor(self)
            name = self._extract_name(serialized, **kwargs)
            self._start_span(run_id, parent_run_id, name, SpanKind.CLIENT, SpanType.GENERATION)

            # Initialize TTFT tracking (None means waiting for first token)
            s.first_token_times[str(run_id)] = None

            if str(run_id) in s.runs:
                span = s.runs[str(run_id)]["span"]
                span.set_attribute(self.ATTR_OPERATION_NAME, "chat")

                # Provider name (try to extract real provider, fallback to "langchain")
                provider = self._extract_provider_name(serialized, kwargs) or "langchain"
                span.set_attribute(self.ATTR_LLM_PROVIDER_NAME, provider)

                # Serialize messages (OTEL standard)
                flat_messages = [msg for batch in messages for msg in batch]
                serialized_messages = [self._serialize_for_json(msg) for msg in flat_messages]
                messages_json = json.dumps(serialized_messages, default=str)
                span.set_attribute(self.ATTR_LLM_INPUT_MESSAGES, messages_json)

                # Extract model name and parameters from serialized kwargs
                if "kwargs" in serialized:
                    model_kwargs = serialized["kwargs"]
                    model_name = model_kwargs.get("model_name") or model_kwargs.get("model")
                    if model_name:
                        span.set_attribute(self.ATTR_LLM_MODEL, model_name)
                        s.runs[str(run_id)]["cached_model_name"] = model_name
                    self._set_model_parameters(span, model_kwargs)

                # Include tools if present (OTEL standard)
                tools = kwargs.get("invocation_params", {}).get("tools")
                if tools:
                    tools_json = json.dumps(tools, default=str)
                    span.set_attribute(self.ATTR_TOOL_DEFINITIONS, tools_json)

        except Exception as e:
            logger.error(f"Error in on_chat_model_start: {e}")

    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle LLM end event."""
        try:
            # DEBUG: Token usage investigation for structured output
            _gen = response.generations[0][0] if response.generations and response.generations[0] else None
            _msg = getattr(_gen, "message", None) if _gen else None
            _usage_meta = getattr(_msg, "usage_metadata", "NO_ATTR") if _msg else "NO_MESSAGE"
            _streaming_usage = self._state.runs.get(str(run_id), {}).get("streaming_usage") if hasattr(self, "_state") else "NO_STATE"
            logger.warning(
                f"[BEACON DEBUG] on_llm_end run_id={run_id} "
                f"llm_output={response.llm_output} "
                f"usage_metadata={_usage_meta} "
                f"streaming_usage={_streaming_usage}"
            )
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            s = _StateAccessor(self)
            run_id_str = str(run_id)
            if run_id_str in s.runs:
                span = s.runs[run_id_str]["span"]

                # Extract response
                if response.generations and response.generations[0]:
                    gen = response.generations[0][0]
                    if hasattr(gen, "message"):
                        serialized_output = self._serialize_for_json(gen.message)
                        completion_json = json.dumps(serialized_output, default=str)
                        span.set_attribute(self.ATTR_LLM_OUTPUT_MESSAGES, completion_json)
                    elif hasattr(gen, "text"):
                        completion_json = json.dumps({"text": gen.text}, default=str)
                        span.set_attribute(self.ATTR_LLM_OUTPUT_MESSAGES, completion_json)

                    # Extract finish reason from generation_info
                    gen_info = getattr(gen, "generation_info", None)
                    if isinstance(gen_info, dict):
                        reason = gen_info.get("finish_reason") or gen_info.get("stop_reason")
                        if reason:
                            span.set_attribute(self.ATTR_RESPONSE_FINISH_REASONS, json.dumps([reason]))

                    # Extract response ID and finish reason from response_metadata
                    if hasattr(gen, "message") and hasattr(gen.message, "response_metadata"):
                        resp_meta = gen.message.response_metadata
                        if resp_meta:
                            resp_id = resp_meta.get("id") or resp_meta.get("response_id")
                            if resp_id:
                                span.set_attribute(self.ATTR_RESPONSE_ID, resp_id)
                            if not span.attributes.get(self.ATTR_RESPONSE_FINISH_REASONS):
                                fr = resp_meta.get("finish_reason") or resp_meta.get("stop_reason")
                                if fr:
                                    span.set_attribute(self.ATTR_RESPONSE_FINISH_REASONS, json.dumps([fr]))

                # Extract usage and model name - check multiple locations
                usage = None
                model_name = None

                # First try: llm_output (non-streaming)
                if response.llm_output:
                    usage = response.llm_output.get("token_usage") or response.llm_output.get("usage")
                    model_name = response.llm_output.get("model_name") or response.llm_output.get("model")
                    # Extract response ID from llm_output
                    if not span.attributes.get(self.ATTR_RESPONSE_ID):
                        resp_id = response.llm_output.get("id")
                        if resp_id:
                            span.set_attribute(self.ATTR_RESPONSE_ID, resp_id)

                # Prefer usage_metadata (TypedDict with input_token_details/output_token_details)
                # over raw llm_output dicts which lack detail breakdowns
                if response.generations and response.generations[0]:
                    gen = response.generations[0][0]
                    if hasattr(gen, "message") and hasattr(gen.message, "usage_metadata"):
                        if gen.message.usage_metadata:
                            usage = gen.message.usage_metadata

                # Fallback: usage captured from streaming chunks (on_llm_new_token)
                if not usage and run_id_str in s.runs:
                    usage = s.runs[run_id_str].get('streaming_usage')

                # Try to get model from response_metadata (streaming)
                if not model_name and response.generations and response.generations[0]:
                    gen = response.generations[0][0]
                    if hasattr(gen, "message") and hasattr(gen.message, "response_metadata"):
                        response_metadata = gen.message.response_metadata
                        if response_metadata:
                            model_name = response_metadata.get("model_name") or response_metadata.get("model")

                # Fallback: use cached model from on_chat_model_start/on_llm_start
                if not model_name:
                    model_name = s.runs[run_id_str].get("cached_model_name")

                if model_name:
                    span.set_attribute(self.ATTR_LLM_MODEL, model_name)
                    span.set_attribute(self.ATTR_RESPONSE_MODEL, model_name)

                if usage:
                    self._set_token_usage(span, usage)

                # Calculate and set TTFT for streaming calls
                first_token_time = s.first_token_times.pop(run_id_str, None)
                if first_token_time is not None:
                    # Streaming call - calculate TTFT
                    start_perf_time = s.runs[run_id_str].get("start_perf_time")
                    if start_perf_time is not None:
                        ttft_ms = (first_token_time - start_perf_time) * 1000
                        span.set_attribute(self.ATTR_TTFT_MS, ttft_ms)
                        span.set_attribute(self.ATTR_STREAMING, True)
                else:
                    # Clean up tracking state even if no first token was received
                    s.first_token_times.pop(run_id_str, None)

            self._end_span(run_id)

        except Exception as e:
            logger.error(f"Error in on_llm_end: {e}")

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle LLM error event."""
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            # Clean up TTFT tracking state
            _StateAccessor(self).first_token_times.pop(str(run_id), None)
            self._end_span(run_id, error)
        except Exception as e:
            logger.error(f"Error in on_llm_error: {e}")

    def on_llm_new_token(
        self,
        token: str,
        *,
        chunk: Any = None,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle new token event for streaming LLM calls.

        This callback captures the time-to-first-token (TTFT) metric by
        recording the timestamp when the first token is received.

        It also captures usage information from streaming chunks (OpenAI sends
        usage in the final chunk).

        Args:
            token: The new token string
            chunk: Optional generation chunk object
            run_id: Unique run identifier
            parent_run_id: Parent run identifier (for nesting)
            **kwargs: Additional arguments
        """
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            s = _StateAccessor(self)
            run_id_str = str(run_id)

            # Only record the first token time
            if run_id_str in s.first_token_times and s.first_token_times[run_id_str] is None:
                s.first_token_times[run_id_str] = time.perf_counter()

            if run_id_str in s.runs:
                run_data = s.runs[run_id_str]
                # Capture usage from streaming chunks (OpenAI sends usage in final chunk)
                if chunk and hasattr(chunk, 'usage_metadata') and chunk.usage_metadata:
                    run_data['streaming_usage'] = chunk.usage_metadata

                # Accumulate streamed content for partial output on cancellation
                # token can be str (most providers) or list (Anthropic content blocks)
                if isinstance(token, str) and token:
                    if 'streaming_content' not in run_data:
                        run_data['streaming_content'] = []
                    run_data['streaming_content'].append(token)
                elif isinstance(token, list):
                    # Anthropic content blocks: [{"type": "text", "text": "..."}]
                    # or thinking blocks: [{"type": "thinking", "thinking": "..."}]
                    for block in token:
                        if not isinstance(block, dict):
                            continue
                        block_type = block.get("type")
                        if block_type == "text":
                            text_val = block.get("text", "")
                            if text_val:
                                if 'streaming_content' not in run_data:
                                    run_data['streaming_content'] = []
                                run_data['streaming_content'].append(text_val)
                        elif block_type == "thinking":
                            thinking_val = block.get("thinking", "")
                            if thinking_val:
                                if 'streaming_thinking' not in run_data:
                                    run_data['streaming_thinking'] = []
                                run_data['streaming_thinking'].append(thinking_val)

        except Exception as e:
            logger.debug(f"Error in on_llm_new_token: {e}")

    # Tool callbacks

    def on_tool_start(
        self,
        serialized: dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        inputs: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle tool start event."""
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            s = _StateAccessor(self)
            name = self._extract_name(serialized, **kwargs)
            self._start_span(run_id, parent_run_id, name, SpanKind.INTERNAL, SpanType.TOOL)

            # Track active tool for nested graph parenting
            s.active_tool_run_ids.append(str(run_id))

            if str(run_id) in s.runs:
                span = s.runs[str(run_id)]["span"]
                span.set_attribute(self.ATTR_OPERATION_NAME, "execute_tool")
                span.set_attribute(self.ATTR_TOOL_NAME, name)
                span.set_attribute(self.ATTR_TOOL_TYPE, "function")
                input_data = inputs if inputs else input_str
                serialized_input = self._serialize_for_json(input_data)
                input_json = json.dumps(serialized_input, default=str)
                span.set_attribute(self.ATTR_TOOL_CALL_ARGUMENTS, input_json)

        except Exception as e:
            logger.error(f"Error in on_tool_start: {e}")

    def on_tool_end(
        self,
        output: Any,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle tool end event."""
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            s = _StateAccessor(self)
            # Remove from active tools stack
            run_id_str = str(run_id)
            if run_id_str in s.active_tool_run_ids:
                s.active_tool_run_ids.remove(run_id_str)
            self._claimed_tool_run_ids.discard(run_id_str)

            if run_id_str in s.runs:
                span = s.runs[run_id_str]["span"]
                serialized_output = self._serialize_for_json(output)
                output_str = json.dumps(serialized_output, default=str)
                span.set_attribute(self.ATTR_TOOL_CALL_RESULT, output_str)

            self._end_span(run_id)

        except Exception as e:
            logger.error(f"Error in on_tool_end: {e}")

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle tool error event."""
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            # Remove from active tools stack
            s = _StateAccessor(self)
            run_id_str = str(run_id)
            if run_id_str in s.active_tool_run_ids:
                s.active_tool_run_ids.remove(run_id_str)
            self._claimed_tool_run_ids.discard(run_id_str)

            self._end_span(run_id, error)
        except Exception as e:
            logger.error(f"Error in on_tool_error: {e}")

    # Retriever callbacks

    def on_retriever_start(
        self,
        serialized: dict[str, Any],
        query: str,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle retriever start event."""
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            s = _StateAccessor(self)
            name = self._extract_name(serialized, **kwargs)
            self._start_span(run_id, parent_run_id, name, SpanKind.INTERNAL, SpanType.RETRIEVAL)

            if str(run_id) in s.runs:
                span = s.runs[str(run_id)]["span"]
                span.set_attribute("retriever.query", query)

        except Exception as e:
            logger.error(f"Error in on_retriever_start: {e}")

    def on_retriever_end(
        self,
        documents: Sequence[Document],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle retriever end event."""
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            s = _StateAccessor(self)
            if str(run_id) in s.runs:
                span = s.runs[str(run_id)]["span"]
                span.set_attribute("retriever.document_count", len(documents))

                # Serialize documents
                serialized_docs = [self._serialize_for_json(doc) for doc in documents]
                docs_str = json.dumps(serialized_docs, default=str)
                span.set_attribute("retriever.documents", docs_str)

            self._end_span(run_id)

        except Exception as e:
            logger.error(f"Error in on_retriever_end: {e}")

    def on_retriever_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle retriever error event."""
        try:
            self._ensure_auto_handoff_context(run_id, parent_run_id)
            self._end_span(run_id, error)
        except Exception as e:
            logger.error(f"Error in on_retriever_error: {e}")

    # Agent callbacks

    def on_agent_action(
        self,
        action: AgentAction,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle agent action event."""
        try:
            s = _StateAccessor(self)
            parent_id = str(parent_run_id) if parent_run_id else None
            if parent_id and parent_id in s.runs:
                span = s.runs[parent_id]["span"]
                action_data = {
                    "tool": action.tool,
                    "tool_input": action.tool_input,
                    "log": action.log if hasattr(action, "log") else None,
                }
                span.set_attribute(f"agent.action.{str(run_id)[:8]}", json.dumps(action_data, default=str))

        except Exception as e:
            logger.error(f"Error in on_agent_action: {e}")

    def on_agent_finish(
        self,
        finish: AgentFinish,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        """Handle agent finish event."""
        try:
            s = _StateAccessor(self)
            parent_id = str(parent_run_id) if parent_run_id else None
            if parent_id and parent_id in s.runs:
                span = s.runs[parent_id]["span"]
                finish_data = {
                    "return_values": finish.return_values,
                    "log": finish.log if hasattr(finish, "log") else None,
                }
                span.set_attribute("agent.finish", json.dumps(finish_data, default=str))

        except Exception as e:
            logger.error(f"Error in on_agent_finish: {e}")


class BeaconLangGraphConfig:
    """LangGraph config builder with automatic Beacon trace continuation.

    Reads checkpoint metadata to detect interrupt state and continue
    traces across interrupt/resume cycles.  Builds LangGraph configs
    with beacon metadata for checkpoint persistence.

    The contained :attr:`handler` is a plain
    :class:`BeaconLangGraphHandler` that can be used independently.

    Args:
        graph: Compiled LangGraph graph (must have a checkpointer).
        thread_id: LangGraph thread ID for checkpointing.
        agent_name: Human-readable agent name for spans.
        session_id: Optional session ID for conversation grouping.
        user_email: User email for attribution (optional).
        user_name: User display name for attribution (optional).
        user_id: User ID for attribution (optional).
        environment: Environment name (e.g., "production", "staging").
        metadata: Custom metadata dict to add to all spans.
        autodetect_handoffs: When True, nested graph invocations
            under tools that have LangGraph metadata are automatically
            treated as handoffs (separate traces with HANDOFF spans).
        governance: Optional :class:`~lumenova_beacon.governance.config.GovernanceConfig`
            to enable governance policy enforcement.  When provided, a
            :class:`~lumenova_beacon.governance.integrations.langchain.BeaconLangGraphGovernanceHandler`
            is created automatically and included in the callbacks.
            Shared parameters (``agent_name``, ``session_id``, ``environment``,
            ``metadata``) are propagated from this config.
        record_stacktrace: Capture full stack traces on exceptions
            (default: inherits from client config).

    Example (sync checkpointer)::

        beacon = BeaconLangGraphConfig(
            graph=agent.graph,
            thread_id=thread_id,
            agent_name="planner",
        )
        config = beacon.get_config()

    Example (with governance)::

        from lumenova_beacon.governance import GovernanceConfig

        beacon = BeaconLangGraphConfig(
            graph=agent.graph,
            thread_id=thread_id,
            agent_name="planner",
            governance=GovernanceConfig(stack_ids=["my-stack"]),
        )
        config = beacon.get_config()

    Example (async checkpointer, e.g. AsyncPostgresSaver)::

        beacon = BeaconLangGraphConfig(
            graph=agent.graph,
            thread_id=thread_id,
            agent_name="planner",
        )
        config = await beacon.aget_config()
    """


    def __init__(
        self,
        graph: Any,
        thread_id: str,
        agent_name: str | None = None,
        session_id: str | None = None,
        user_email: str | None = None,
        user_name: str | None = None,
        user_id: str | None = None,
        environment: str | None = None,
        metadata: dict[str, Any] | None = None,
        autodetect_handoffs: bool = False,
        governance: GovernanceConfig | None = None,
        record_stacktrace: bool | None = None,
        **kwargs: Any,
    ):
        self._graph = graph
        self._thread_id = thread_id
        self._user_session_id = session_id  # preserve original to detect user override

        # Pre-generate default IDs (may be overridden by checkpoint resolution)
        if session_id is None:
            import uuid as _uuid
            session_id = str(_uuid.uuid4())

        self._handler = BeaconLangGraphHandler(
            agent_name=agent_name,
            session_id=session_id,
            user_email=user_email,
            user_name=user_name,
            user_id=user_id,
            environment=environment,
            metadata=metadata,
            autodetect_handoffs=autodetect_handoffs,
            record_stacktrace=record_stacktrace,
            **kwargs,
        )
        self._handler._current_trace_id = generate_trace_id()
        self._handler._root_span_id = generate_span_id()
        self._handler._is_resume = False
        self._resolved = False

        # Governance integration — create handler when GovernanceConfig is provided.
        if governance is not None:
            from lumenova_beacon.governance.integrations.langchain import (
                BeaconLangGraphGovernanceHandler,
            )

            self._governance_handler = BeaconLangGraphGovernanceHandler(
                stack_ids=governance.stack_ids,
                tags=governance.tags,
                fail_open=governance.fail_open,
                enabled_hooks=governance.enabled_hooks,
                result_parser=governance.result_parser,
                timeout=governance.timeout,
                debug=governance.debug,
                on_violation=governance.on_violation,
                max_violations=governance.max_violations,
                agent_id=governance.agent_id,
                agent_name=agent_name,
                session_id=session_id,
                environment=environment,
                metadata=metadata,
            )
        else:
            self._governance_handler = None

    @property
    def handler(self) -> BeaconLangGraphHandler:
        """The callback handler instance."""
        return self._handler

    @property
    def trace_id(self) -> str | None:
        """Current trace ID."""
        return self._handler.trace_id

    @property
    def root_span_id(self) -> str | None:
        """Root span ID for resume continuation."""
        return self._handler.root_span_id

    @property
    def is_resume(self) -> bool:
        """Whether this is a resumed trace."""
        return self._handler._is_resume

    @property
    def governance_handler(self) -> Any:
        """The governance handler instance, or ``None`` if not configured."""
        return self._governance_handler

    def mark_interrupt(
        self,
        value: str | None = None,
        interrupt_id: str | None = None,
    ) -> None:
        """Mark the current trace as interrupted.

        Delegates to :meth:`BeaconLangGraphHandler.mark_interrupt` so the
        root span gets interrupt attributes.  Call this on user
        cancellation or any external interrupt that the callback handler
        cannot detect on its own.

        Args:
            value: Human-readable interrupt description.
            interrupt_id: Checkpoint ID for linking the subsequent resume.
        """
        self._handler.mark_interrupt(value=value, interrupt_id=interrupt_id)

    @property
    def checkpoint_metadata(self) -> dict[str, str]:
        """Beacon keys for LangGraph checkpoint persistence.

        Spread into ``config["metadata"]`` so LangGraph persists them
        with the checkpoint (only flat primitives survive serialization).
        """
        result: dict[str, str] = {}
        if self._handler._current_trace_id:
            result["beacon_trace_id"] = self._handler._current_trace_id
        if self._handler._root_span_id:
            result["beacon_root_span_id"] = self._handler._root_span_id
        if self._handler._session_id:
            result["beacon_session_id"] = self._handler._session_id
        return result

    def _apply_checkpoint(self, resolved: dict[str, Any] | None) -> None:
        """Apply resolved checkpoint data to handler state."""
        if resolved:
            if resolved["trace_id"]:
                self._handler._current_trace_id = resolved["trace_id"]
            if resolved["root_span_id"]:
                self._handler._root_span_id = resolved["root_span_id"]
            self._handler._is_resume = resolved["is_resume"]
            if resolved.get("session_id") and self._user_session_id is None:
                self._handler._session_id = resolved["session_id"]
                if self._governance_handler is not None:
                    self._governance_handler._session_id = resolved["session_id"]
        self._resolved = True

    def _build_config(
        self,
        configurable: dict[str, Any] | None = None,
        metadata: dict[str, Any] | None = None,
        tags: list[str] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Build the RunnableConfig dict (shared by sync/async paths)."""
        callbacks: list[Any] = [self._handler]
        if self._governance_handler is not None:
            callbacks.append(self._governance_handler)
        config: dict[str, Any] = {
            "configurable": {"thread_id": self._thread_id},
            "metadata": dict(self.checkpoint_metadata),
            "callbacks": callbacks,
        }
        if configurable:
            config["configurable"].update(configurable)
        if metadata:
            config["metadata"].update(metadata)
        if tags:
            config["tags"] = tags
        if kwargs:
            config.update(kwargs)
        return config

    def get_config(
        self,
        *,
        configurable: dict[str, Any] | None = None,
        metadata: dict[str, Any] | None = None,
        tags: list[str] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Build a LangGraph config with beacon trace persistence.

        Resolves checkpoint state synchronously on first call.
        Use :meth:`aget_config` with async checkpointers
        (e.g. ``AsyncPostgresSaver``).

        Args:
            configurable: Additional configurable keys (merged with
                ``thread_id``).
            metadata: Additional metadata keys (merged with beacon IDs).
            tags: Tags for the run.
            **kwargs: Any other ``RunnableConfig`` keys.

        Returns:
            Complete ``RunnableConfig`` dict.

        Example::

            config = beacon.get_config(
                tags=["production"],
                configurable={"recursion_limit": 50},
            )
        """
        if not self._resolved:
            resolved = self._resolve_from_checkpoint(self._graph, self._thread_id)
            self._apply_checkpoint(resolved)
        return self._build_config(configurable=configurable, metadata=metadata, tags=tags, **kwargs)

    async def aget_config(
        self,
        *,
        configurable: dict[str, Any] | None = None,
        metadata: dict[str, Any] | None = None,
        tags: list[str] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Build a LangGraph config with beacon trace persistence (async).

        Resolves checkpoint state asynchronously on first call.
        Required for async checkpointers (e.g. ``AsyncPostgresSaver``).

        Args:
            configurable: Additional configurable keys (merged with
                ``thread_id``).
            metadata: Additional metadata keys (merged with beacon IDs).
            tags: Tags for the run.
            **kwargs: Any other ``RunnableConfig`` keys.

        Returns:
            Complete ``RunnableConfig`` dict.

        Example::

            config = await beacon.aget_config(
                tags=["production"],
                configurable={"recursion_limit": 50},
            )
        """
        if not self._resolved:
            resolved = await self._aresolve_from_checkpoint(self._graph, self._thread_id)
            self._apply_checkpoint(resolved)
        return self._build_config(configurable=configurable, metadata=metadata, tags=tags, **kwargs)

    @property
    def langgraph_config(self) -> dict[str, Any]:
        """Complete LangGraph config (alias for ``get_config()``)."""
        return self.get_config()

    @staticmethod
    def _resolve_from_checkpoint(graph: Any, thread_id: str) -> dict[str, Any] | None:
        """Read checkpoint metadata to resolve trace continuation state.

        Uses the sync ``graph.get_state()`` which works with any
        checkpointer (MemorySaver, PostgresSaver, etc.).
        """
        try:
            state = graph.get_state({"configurable": {"thread_id": thread_id}})

            if not state or not state.metadata:
                return None

            stored_trace_id = state.metadata.get("beacon_trace_id")
            stored_root_span_id = state.metadata.get("beacon_root_span_id")
            stored_session_id = state.metadata.get("beacon_session_id")
            is_interrupted = bool(state.next) or bool(getattr(state, "interrupts", ()))

            if stored_trace_id and is_interrupted:
                return {
                    "trace_id": stored_trace_id,
                    "root_span_id": stored_root_span_id,
                    "is_resume": True,
                    "session_id": stored_session_id,
                }

            if stored_session_id:
                return {
                    "trace_id": None,
                    "root_span_id": None,
                    "is_resume": False,
                    "session_id": stored_session_id,
                }
        except Exception:
            pass
        return None

    @staticmethod
    async def _aresolve_from_checkpoint(graph: Any, thread_id: str) -> dict[str, Any] | None:
        """Async version of checkpoint resolution for async checkpointers.

        Uses ``await graph.aget_state()`` which is required for async
        checkpointers (e.g. ``AsyncPostgresSaver``).
        """
        try:
            state = await graph.aget_state({"configurable": {"thread_id": thread_id}})

            if not state or not state.metadata:
                return None

            stored_trace_id = state.metadata.get("beacon_trace_id")
            stored_root_span_id = state.metadata.get("beacon_root_span_id")
            stored_session_id = state.metadata.get("beacon_session_id")
            is_interrupted = bool(state.next) or bool(getattr(state, "interrupts", ()))

            if stored_trace_id and is_interrupted:
                return {
                    "trace_id": stored_trace_id,
                    "root_span_id": stored_root_span_id,
                    "is_resume": True,
                    "session_id": stored_session_id,
                }

            if stored_session_id:
                return {
                    "trace_id": None,
                    "root_span_id": None,
                    "is_resume": False,
                    "session_id": stored_session_id,
                }
        except Exception:
            pass
        return None


# Deprecated aliases — old names emit a DeprecationWarning on access.
_DEPRECATED_NAMES = {
    "BeaconCallbackHandler": "BeaconLangGraphHandler",
    "LangGraphBeaconConfig": "BeaconLangGraphConfig",
}


def __getattr__(name: str):
    if name in _DEPRECATED_NAMES:
        import warnings

        new_name = _DEPRECATED_NAMES[name]
        warnings.warn(
            f"{name} is deprecated, use {new_name} instead",
            DeprecationWarning,
            stacklevel=2,
        )
        return globals()[new_name]
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
