app/observability/operation.py

"""The single operation boundary: one span, bounded metrics, one log, one row.

Every instrumented call site in later tasks wraps its work in
:func:`observe_operation`. The context manager reads the clock **once** and
feeds that single duration to the span, the duration histogram, and the local
observation -- there is no second timer anywhere downstream. On completion it:

1. records one monotonic start time;
2. creates an OTel span when mode is ``local`` or ``full``;
3. attaches sanitized trace attributes;
4. exposes ``set()`` / ``event()`` / ``mark_error()`` / ``mark_terminal()``;
5. records the duration plus bounded RAG stage/measurement histograms;
6. emits a structured correlated log for errors/fallbacks;
7. writes the same completed :class:`OperationObservation` to the local store;
8. suppresses and counts every exporter/store/SDK failure so a telemetry
   problem can never change or block a product result.

``store`` and ``exporter`` are injectable for tests and specialised callers;
when omitted they resolve from the initialized global state. An explicitly
passed ``store`` is always honoured, even before/without global init.
"""

from __future__ import annotations

import logging
import time
from contextlib import contextmanager
from typing import Any, Iterator, Protocol

from opentelemetry import trace
from opentelemetry.trace import Span, Status, StatusCode

from app.observability import bootstrap
from app.observability.contracts import (
    METRIC_ATTRIBUTE_KEYS,
    ExperimentIdentity,
    OperationObservation,
    RetrievalOutcome,
    metric_attributes,
)
from app.observability.logging import emit_correlated_log, trace_context_ids
from app.observability.rag_signals import normalize_measure, normalize_stage
from app.observability.sanitize import (
    REDACTED,
    is_forbidden_key,
    sanitize_attribute_value,
    sanitize_attributes,
    sanitize_exception,
)

logger = logging.getLogger(__name__)

# Statuses that count as a failure/fallback for logging purposes.
_ERROR_STATUSES = frozenset({"error", "error_fallback", "degraded"})


class _Exporter(Protocol):
    def export(self, observation: OperationObservation) -> Any: ...


def _guard(fn, *, count: bool = True) -> None:
    """Run a telemetry side-effect, swallowing (and optionally counting) failure."""
    try:
        fn()
    except Exception:
        if count:
            bootstrap.note_suppressed_failure()
        logger.debug("observability side-effect suppressed", exc_info=True)


class Operation:
    """Handle yielded by :func:`observe_operation`.

    Product code calls ``set``/``event``/``add_stage``/``add_count`` and the
    ``set_experiment``/``set_retrieval``/``mark_*`` helpers; everything is
    buffered and sanitized, then flushed once at ``__exit__``.
    """

    def __init__(
        self,
        name: str,
        *,
        subsystem: str | None,
        consumer: str | None,
        evaluation_run: bool,
    ) -> None:
        self.name = name
        self.subsystem = subsystem
        self.consumer = consumer
        self.evaluation_run = evaluation_run
        self._span: Span | None = None
        self._attributes: dict[str, Any] = {}
        self._stages: dict[str, float] = {}
        self._counts: dict[str, int | float] = {}
        self._exportable_stages: dict[str, float] = {}
        self._exportable_measures: dict[str, float] = {}
        self._experiment: ExperimentIdentity | None = None
        self._retrieval: RetrievalOutcome | None = None
        self._status: str | None = None
        self._errored = False
        self._exception_recorded = False
        self._recorded_exception: BaseException | None = None

    # -- span binding -------------------------------------------------

    def _bind_span(self, span: Span | None) -> None:
        self._span = span
        if span is None:
            return
        self._span_set("operation", self.name)
        if self.subsystem:
            self._span_set("subsystem", self.subsystem)
        if self.consumer:
            self._span_set("consumer", self.consumer)
        if self.evaluation_run:
            self._span_set("evaluation_run", True)

    def _span_set(self, key: str, value: Any) -> None:
        if self._span is None:
            return
        _guard(lambda: self._span.set_attribute(key, value), count=False)

    # -- public API -------------------------------------------------

    def set(self, key: str, value: Any) -> None:
        """Attach one sanitized attribute to the observation and span."""
        skey = str(key)
        if is_forbidden_key(skey):
            sanitized: Any = REDACTED
        else:
            sanitized = sanitize_attribute_value(value)
            if sanitized is None:
                return
        self._attributes[skey] = sanitized
        self._span_set(skey, sanitized)

    def event(self, name: str, attributes: dict[str, Any] | None = None) -> None:
        """Record a timestamped span event (a discrete in-operation moment)."""
        clean = sanitize_attributes(attributes or {})
        if self._attributes.get("diagnostic.smoke") is True:
            clean["diagnostic.smoke"] = True
        if self._span is not None:
            _guard(lambda: self._span.add_event(name, attributes=clean), count=False)

    def add_stage(self, name: str, duration_ms: float) -> None:
        stage = normalize_stage(name)
        self._stages[str(name)] = float(duration_ms)
        if stage:
            self._exportable_stages[stage] = float(duration_ms)
            self._span_set(f"rag.stage.{stage}.ms", float(duration_ms))

    @property
    def stage_durations_ms(self) -> dict[str, float]:
        """Return a read-only snapshot of recorded stage durations."""

        return dict(self._stages)

    def add_count(self, name: str, value: int | float) -> None:
        measure = normalize_measure(name)
        self._counts[str(name)] = value
        if measure:
            self._exportable_measures[measure] = float(value)
            self._span_set(f"rag.measure.{measure}", float(value))

    def set_experiment(self, experiment: ExperimentIdentity) -> None:
        self._experiment = experiment
        for key, value in {
            "experiment.id": experiment.experiment_id,
            "experiment.run_id": experiment.run_id,
            "pipeline.version": experiment.pipeline_version,
            "query.category": experiment.query_category,
            "cache.state": experiment.cache_state,
            "telemetry.mode": experiment.telemetry_mode,
            "cold_start": experiment.cold_start,
        }.items():
            if value not in (None, ""):
                self._span_set(key, value)

    def set_retrieval(self, retrieval: RetrievalOutcome) -> None:
        self._retrieval = retrieval

    def mark_error(self, error_type: str | None = None) -> None:
        """Flag this operation as failed without raising an exception."""
        self._errored = True
        self._status = "error"
        if error_type:
            self._attributes["error.type"] = sanitize_attribute_value(str(error_type))

    def mark_terminal(self, status: str) -> None:
        """Set an explicit terminal status (e.g. ``success_empty``)."""
        self._status = str(status)
        if status in _ERROR_STATUSES:
            self._errored = True

    def record_exception(
        self,
        exc: BaseException,
        *,
        escaped: bool,
        stage: str | None = None,
        category: str | None = None,
    ) -> None:
        """Add one sanitized semantic exception event to this operation.

        This intentionally does not call OpenTelemetry's default exception
        recorder: its traceback payload is not safe for the research path.
        """
        if escaped:
            self._errored = True
            self._status = "error"
        try:
            exception_type, message = sanitize_exception(exc)
            normalized_stage = normalize_stage(stage) if stage is not None else None
            event: dict[str, Any] = {
                "exception.type": exception_type,
                "exception.message": message,
                "exception.escaped": bool(escaped),
            }
            if normalized_stage:
                event["rag.stage"] = normalized_stage
            if category:
                event["error.category"] = sanitize_attribute_value(category)
            self.event("exception", event)
            if escaped:
                self._attributes.setdefault("error.type", exception_type)
                self._attributes.setdefault("error.message", message)
                self._span_set("error.type", self._attributes["error.type"])
                self._span_set("error.message", self._attributes["error.message"])
        except BaseException:
            # Observability must never replace the product exception, including
            # if a third-party span implementation is itself defective.
            logger.debug("exception observability suppressed")
        finally:
            self._exception_recorded = True
            self._recorded_exception = exc

    # -- finalisation -------------------------------------------------

    def _resolve_status(self) -> str:
        if self._status is not None:
            return self._status
        return "error" if self._errored else "success"

    def _metric_dimensions(self, status: str) -> dict[str, Any]:
        dims: dict[str, Any] = {"operation": self.name, "status": status}
        if self.subsystem:
            dims["subsystem"] = self.subsystem
        if self.consumer:
            dims["consumer"] = self.consumer
        dims["evaluation_run"] = self.evaluation_run
        if self._experiment is not None:
            if self._experiment.pipeline_version:
                dims["pipeline.version"] = self._experiment.pipeline_version
            if self._experiment.query_category:
                dims["query.category"] = self._experiment.query_category
            if self._experiment.cache_state:
                dims["cache.state"] = self._experiment.cache_state
        for key in METRIC_ATTRIBUTE_KEYS:
            if key in self._attributes:
                dims.setdefault(key, self._attributes[key])
        return metric_attributes(dims)

    def _build_observation(self, status: str, duration_ms: float | None) -> OperationObservation:
        trace_id, _span_id = trace_context_ids(self._span)
        return OperationObservation(
            operation=self.name,
            status=status,
            subsystem=self.subsystem,
            consumer=self.consumer,
            duration_ms=duration_ms,
            stage_durations_ms=dict(self._stages),
            counts=dict(self._counts),
            trace_id=trace_id,
            experiment=self._experiment,
            retrieval=self._retrieval,
            evaluation_run=self.evaluation_run,
            attributes=dict(self._attributes),
        )

    def _finalize_span(self, status: str) -> None:
        if self._span is None:
            return
        self._span_set("status", status)
        if self._errored or status in _ERROR_STATUSES:
            _guard(lambda: self._span.set_status(Status(StatusCode.ERROR)), count=False)
        else:
            _guard(lambda: self._span.set_status(Status(StatusCode.OK)), count=False)


@contextmanager
def observe_operation(
    name: str,
    attributes: dict[str, Any] | None = None,
    *,
    store: Any = None,
    exporter: _Exporter | None = None,
    subsystem: str | None = None,
    consumer: str | None = None,
    evaluation_run: bool = False,
    record_escaped_exceptions: bool = True,
) -> Iterator[Operation]:
    """Wrap one product operation in the canonical observability boundary.

    Always yields a valid :class:`Operation` so call sites never need
    mode-conditional code. In ``disabled`` mode with no explicitly injected
    store this is a plain timing no-op: no span, no metric, no store write.
    """
    state = bootstrap.get_state()
    resolved_store = store if store is not None else state.local_store
    tracer = state.tracer
    histogram = state.duration_histogram
    stage_histogram = state.stage_duration_histogram
    measurement_histogram = state.rag_measurement_histogram

    op = Operation(
        name,
        subsystem=subsystem,
        consumer=consumer,
        evaluation_run=evaluation_run,
    )

    # (1) one monotonic clock read feeds span, metric, and observation.
    start = time.monotonic()

    # Span setup is guarded: a tracer that fails to start (or enter) a span
    # must degrade to a no-span operation, never propagate into the caller's
    # `with` block. span_cm is only bound once __enter__ has succeeded, so the
    # finally-block never tries to __exit__ a half-started context manager.
    span_cm: Any = None
    span: Span | None = None

    if tracer is not None:
        def _start_span() -> None:
            nonlocal span_cm, span
            cm = tracer.start_as_current_span(name)
            span = cm.__enter__()
            span_cm = cm

        _guard(_start_span, count=False)

    op._bind_span(span)

    if attributes:
        for key, value in attributes.items():
            op.set(key, value)

    exc_to_raise: BaseException | None = None
    try:
        yield op
    except BaseException as exc:  # noqa: BLE001 - re-raised below after recording
        if record_escaped_exceptions:
            if not op._exception_recorded or op._recorded_exception is not exc:
                op.record_exception(exc, escaped=True)
        else:
            # A nested technical boundary (currently a RAG stage) does not
            # own the product fallback decision. Mark its span failed, but
            # defer the single semantic handled/escaped exception event to
            # the parent that either applies a fallback or lets the request
            # fail. This prevents a handled child failure being counted as an
            # escaped request exception.
            op.mark_error(type(exc).__name__)
        exc_to_raise = exc
    finally:
        duration_ms = (time.monotonic() - start) * 1000.0
        status = op._resolve_status()

        op._finalize_span(status)

        dims = op._metric_dimensions(status)
        if histogram is not None:
            _guard(lambda: histogram.record(duration_ms, dims))

        if stage_histogram is not None:
            for stage, stage_duration_ms in op._exportable_stages.items():
                stage_dims = {**dims, "stage": stage}
                _guard(
                    lambda stage_duration_ms=stage_duration_ms, stage_dims=stage_dims: (
                        stage_histogram.record(stage_duration_ms, stage_dims)
                    )
                )

        if measurement_histogram is not None:
            for measure, value in op._exportable_measures.items():
                measure_dims = {**dims, "measure": measure}
                _guard(
                    lambda value=value, measure_dims=measure_dims: (
                        measurement_histogram.record(value, measure_dims)
                    )
                )

        if status in _ERROR_STATUSES or (op._retrieval is not None and op._retrieval.retrieval_error):
            trace_id, span_id = trace_context_ids(span)
            pipeline_version = op._experiment.pipeline_version if op._experiment else None
            error_type = op._attributes.get("error.type")
            diagnostic_smoke = (
                True if op._attributes.get("diagnostic.smoke") is True else None
            )
            _guard(
                lambda: emit_correlated_log(
                    f"operation {op.name} finished with status={status}",
                    operation=op.name,
                    status=status,
                    trace_id=trace_id,
                    span_id=span_id,
                    pipeline_version=pipeline_version,
                    consumer=op.consumer,
                    subsystem=op.subsystem,
                    error_type=error_type,
                    diagnostic_smoke=diagnostic_smoke,
                ),
                count=False,
            )

        observation = op._build_observation(status, round(duration_ms, 3))

        if resolved_store is not None:
            _guard(lambda: resolved_store.record(observation))

        if exporter is not None:
            _guard(lambda: exporter.export(observation))

        if span_cm is not None:
            _guard(lambda: span_cm.__exit__(None, None, None), count=False)

    if exc_to_raise is not None:
        raise exc_to_raise


@contextmanager
def observe_stage(
    parent: Operation,
    stage: str,
    *,
    subsystem: str,
    consumer: str | None = None,
    attributes: dict[str, Any] | None = None,
) -> Iterator[Operation]:
    """Observe one registered RAG stage without risking product behaviour.

    Unknown stage names remain useful in local operation observations, but are
    deliberately not emitted as child spans or bounded stage dimensions.
    """
    normalized = normalize_stage(stage)
    if normalized is None:
        started = time.monotonic()
        try:
            yield parent
        finally:
            parent.add_stage(str(stage), (time.monotonic() - started) * 1000.0)
        return

    started = time.monotonic()
    stage_attributes = {**(attributes or {}), "rag.stage": normalized}
    with observe_operation(
        f"rag.stage.{normalized}",
        subsystem=subsystem,
        consumer=consumer,
        attributes=stage_attributes,
        record_escaped_exceptions=False,
    ) as child:
        try:
            yield child
        except BaseException:
            raise
        finally:
            parent.add_stage(normalized, (time.monotonic() - started) * 1000.0)


def record_event(name: str, attributes: dict[str, Any] | None = None) -> None:
    """Record a standalone event on the currently active span, if any.

    A convenience for call sites that want to mark a discrete moment (a cache
    hit, a fallback taken) without holding the :class:`Operation` handle.
    Sanitizes attributes and is a no-op when no recording span is active.
    """
    span = trace.get_current_span()
    ctx = span.get_span_context() if span is not None else None
    if ctx is None or not ctx.is_valid:
        return
    clean = sanitize_attributes(attributes or {})
    _guard(lambda: span.add_event(name, attributes=clean), count=False)