OpenViking/openviking/telemetry/execution.py

# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
# SPDX-License-Identifier: AGPL-3.0
"""Shared helpers for telemetry-wrapped operation execution."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Awaitable, Callable, Generic, Optional, TypeVar

from openviking.observability.context import (
    bind_operation_observability_context,
    reset_operation_observability_context,
)
from openviking_cli.exceptions import InvalidArgumentError
from openviking_cli.utils import get_logger

from .context import bind_telemetry
from .operation import OperationTelemetry, TelemetrySnapshot
from .request import TelemetryRequest, TelemetrySelection, normalize_telemetry_request
from .span_models import OperationSpanAttributes

T = TypeVar("T")
logger = get_logger(__name__)


@dataclass
class TelemetryExecutionResult(Generic[T]):
    """Executed operation result plus telemetry payloads."""

    result: T
    telemetry: Optional[dict[str, Any]]
    selection: TelemetrySelection


def parse_telemetry_selection(telemetry: TelemetryRequest) -> TelemetrySelection:
    """Validate and normalize a telemetry request for public API usage."""
    try:
        return normalize_telemetry_request(telemetry)
    except ValueError as exc:
        raise InvalidArgumentError(str(exc)) from exc


def build_telemetry_payload(
    snapshot: Optional[TelemetrySnapshot],
    selection: TelemetrySelection,
) -> Optional[dict[str, Any]]:
    """Build a telemetry payload from a finished snapshot."""
    if snapshot is None or not selection.include_payload:
        return None
    return snapshot.to_dict(include_summary=selection.include_summary)


def _log_telemetry_summary(snapshot: Optional[TelemetrySnapshot]) -> None:
    if snapshot is None:
        return
    logger.debug(
        "Telemetry summary (id=%s): %s",
        snapshot.telemetry_id,
        snapshot.summary,
    )


def attach_telemetry_payload(
    result: Any,
    telemetry_payload: Optional[dict[str, Any]],
) -> Any:
    """Attach a telemetry payload to a dict result."""
    if telemetry_payload is None:
        return result

    if result is None:
        payload: dict[str, Any] = {}
        payload["telemetry"] = telemetry_payload
        return payload

    if isinstance(result, dict):
        result["telemetry"] = telemetry_payload
        return result

    return result


# Try to import opentelemetry - will be None if not installed
try:
    from opentelemetry import trace as otel_trace
    from opentelemetry.trace import Status, StatusCode
except ImportError:
    otel_trace = None
    Status = None
    StatusCode = None


def _fill_telemetry_summary_to_span(snapshot, span):
    """fill OperationTelemetry summary into span

    Args:
        snapshot: TelemetrySnapshot object
        span: current OTel span
    """
    if snapshot is None or span is None:
        return

    try:
        # Build OperationSpanAttributes from snapshot for consistent span attributes.
        op_attrs = OperationSpanAttributes.from_telemetry_snapshot(
            operation=snapshot.operation,
            telemetry_id=snapshot.telemetry_id,
            status=snapshot.status,
            summary=snapshot.summary,
        )

        # Apply to span attributes (best-effort).
        span.set_attributes(op_attrs.to_otel_attributes())

        # Attach summary as a span event for debugging (best-effort).
        if snapshot.summary:
            span.add_event("telemetry_summary", attributes={"summary": str(snapshot.summary)})
    except Exception:
        logger.debug("failed to fill telemetry summary into span", exc_info=True)



async def run_with_telemetry(
    *,
    operation: str,
    telemetry: TelemetryRequest,
    fn: Callable[[], Awaitable[T]],
    error_status: str = "error",
) -> TelemetryExecutionResult[T]:
    """Execute an async operation with a bound operation-scoped collector."""
    selection = parse_telemetry_selection(telemetry)
    collector = OperationTelemetry(
        operation=operation,
        enabled=True,
    )
    operation_attrs = OperationSpanAttributes(
        operation=operation,
        telemetry_id=collector.telemetry_id,
    )
    operation_token = bind_operation_observability_context(operation_attrs)

    # If OTel is available, create an operation span (best-effort).
    span_context_manager = None
    if otel_trace is not None:
        try:
            tracer = otel_trace.get_tracer(__name__)
            span_context_manager = tracer.start_as_current_span(
                name=operation,
                kind=otel_trace.SpanKind.INTERNAL,
            )
        except Exception:
            logger.debug("failed to start operation span", exc_info=True)
            span_context_manager = None

    try:
        try:
            with bind_telemetry(collector):
                if span_context_manager is not None:
                    with span_context_manager as span:
                        try:
                            span.set_attributes(operation_attrs.to_otel_attributes())
                        except Exception:
                            logger.debug(
                                "failed to set operation span attributes (initial)",
                                exc_info=True,
                            )
                        result = await fn()
                else:
                    result = await fn()
        except Exception as exc:
            collector.set_error(operation, type(exc).__name__, str(exc))
            snapshot = collector.finish(status=error_status)
            _log_telemetry_summary(snapshot)

            # If a span exists, record exception and backfill summary (best-effort).
            if otel_trace is not None:
                try:
                    current_span = otel_trace.get_current_span()
                    if current_span is not None and current_span.is_recording():
                        current_span.record_exception(exc)
                        current_span.set_status(Status(StatusCode.ERROR, description=str(exc)))
                        _fill_telemetry_summary_to_span(snapshot, current_span)
                except Exception:
                    logger.debug(
                        "failed to record exception into operation span",
                        exc_info=True,
                    )

            raise

        snapshot = collector.finish(status="ok")
        _log_telemetry_summary(snapshot)

        # If a span exists, backfill summary (best-effort).
        if otel_trace is not None:
            try:
                current_span = otel_trace.get_current_span()
                if current_span is not None and current_span.is_recording():
                    _fill_telemetry_summary_to_span(snapshot, current_span)
            except Exception:
                logger.debug("failed to backfill telemetry summary into span", exc_info=True)

        try:
            if snapshot is not None:
                from openviking.metrics.datasources.telemetry_bridge import (
                    TelemetryBridgeEventDataSource,
                )

                TelemetryBridgeEventDataSource.record_summary(snapshot.summary)
        except Exception:
            logger.debug("failed to record telemetry summary into metrics bridge", exc_info=True)
        telemetry_payload = build_telemetry_payload(
            snapshot,
            selection,
        )
        return TelemetryExecutionResult(
            result=result,
            telemetry=telemetry_payload,
            selection=selection,
        )
    finally:
        reset_operation_observability_context(operation_token)


__all__ = [
    "TelemetryExecutionResult",
    "attach_telemetry_payload",
    "build_telemetry_payload",
    "parse_telemetry_selection",
    "run_with_telemetry",
]