nemo-relay 0.7.1

// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Optional observability integrations for NeMo Relay Core.

use crate::api::event::EventNormalizationExt;
use serde::{Deserialize, Serialize};

/// Copies a projected OTLP attribute to a second attribute name.
///
/// `key` names the fully-qualified projected attribute and `alias` names the
/// additional attribute to emit with the same typed value.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct OtlpAttributeMapping {
    /// Fully-qualified projected attribute to copy.
    pub key: String,
    /// Additional attribute name receiving the copied value.
    pub alias: String,
}

impl OtlpAttributeMapping {
    /// Creates an attribute mapping.
    pub fn new(key: impl Into<String>, alias: impl Into<String>) -> Self {
        Self {
            key: key.into(),
            alias: alias.into(),
        }
    }
}

#[cfg(test)]
use std::sync::Mutex;

#[cfg(test)]
pub(crate) fn test_mutex() -> &'static Mutex<()> {
    crate::shared_runtime::runtime_owner_test_mutex()
}

pub mod atif;
pub mod atof;
pub(crate) mod manual;
pub(crate) mod openinference;
pub mod otel;
mod otel_genai;
pub mod plugin_component;

/// Export representation for point-in-time mark events.
///
/// Marks remain canonical ATOF events regardless of this setting. Exporters
/// apply the selected projection only when translating those events into a
/// downstream trace format.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum MarkProjection {
    /// Use each exporter’s native handling for marks.
    #[default]
    Inherit,
    /// Force marks into exporter-native trace span events.
    Event,
    /// Render non-excluded marks as zero-duration trace child spans so
    /// trace-tree consumers can display them directly. High-volume
    /// `llm.chunk` marks remain exporter-native events.
    Tool,
}

/// Semantic projection emitted by an OpenTelemetry exporter.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum OpenTelemetryType {
    /// Relay's complete lifecycle projection, including `nemo_relay.*` attributes.
    #[default]
    Full,
    /// OpenTelemetry GenAI semantic conventions.
    GenAi,
    /// OpenInference semantic conventions.
    #[serde(rename = "openinference")]
    OpenInference,
}

/// Default mark names excluded from tool projection because they are emitted
/// at high volume and are better represented as exporter-native events.
/// Return the default mark names excluded from OpenTelemetry projections.
pub fn default_mark_exclude_names() -> Vec<String> {
    vec!["llm.chunk".to_string()]
}

pub(crate) fn relay_trace_id(uuid: uuid::Uuid) -> opentelemetry::trace::TraceId {
    opentelemetry::trace::TraceId::from_bytes(*uuid.as_bytes())
}

pub(crate) fn relay_span_id(uuid: uuid::Uuid) -> opentelemetry::trace::SpanId {
    let mut bytes = [0; 8];
    bytes.copy_from_slice(&uuid.as_bytes()[8..]);
    opentelemetry::trace::SpanId::from_bytes(bytes)
}

pub(crate) fn push_common_optimization_attributes(
    attributes: &mut Vec<opentelemetry::KeyValue>,
    summary: &crate::codec::optimization::LlmOptimizationSummary,
) {
    push_optimization_models_and_tokens(attributes, summary);
    push_optimization_cost(attributes, "baseline", summary.baseline_cost.as_ref());
    push_optimization_cost(attributes, "actual", summary.actual_cost.as_ref());
    push_optimization_savings_and_status(attributes, summary);
    push_optimization_pricing_provenance(attributes, summary);
}

fn push_optimization_models_and_tokens(
    attributes: &mut Vec<opentelemetry::KeyValue>,
    summary: &crate::codec::optimization::LlmOptimizationSummary,
) {
    if let Some(model) = summary.baseline_model.as_ref() {
        attributes.push(opentelemetry::KeyValue::new(
            "nemo_relay.llm.optimization.baseline_model",
            model.model.clone(),
        ));
    }
    if let Some(model) = summary.effective_model.as_ref() {
        attributes.push(opentelemetry::KeyValue::new(
            "nemo_relay.llm.optimization.effective_model",
            model.model.clone(),
        ));
    }
    if let Some(tokens) = summary.tokens_saved.prompt_tokens {
        attributes.push(opentelemetry::KeyValue::new(
            "nemo_relay.llm.optimization.prompt_tokens_saved",
            i64::try_from(tokens).unwrap_or(i64::MAX),
        ));
    }
    if let Some(tokens) = summary.tokens_saved.total_tokens {
        attributes.push(opentelemetry::KeyValue::new(
            "nemo_relay.llm.optimization.total_tokens_saved",
            i64::try_from(tokens).unwrap_or(i64::MAX),
        ));
    }
}

fn push_optimization_cost(
    attributes: &mut Vec<opentelemetry::KeyValue>,
    label: &str,
    cost: Option<&crate::codec::response::CostEstimate>,
) {
    let Some(cost) = cost else {
        return;
    };
    if let Some(total) = cost.total_or_component_sum() {
        attributes.push(opentelemetry::KeyValue::new(
            format!("nemo_relay.llm.optimization.{label}_cost"),
            total,
        ));
    }
    attributes.push(opentelemetry::KeyValue::new(
        format!("nemo_relay.llm.optimization.{label}_cost_currency"),
        cost.currency.clone(),
    ));
    if let Some(source) = cost.pricing_source.as_ref() {
        attributes.push(opentelemetry::KeyValue::new(
            format!("nemo_relay.llm.optimization.{label}_pricing_source"),
            source.clone(),
        ));
    }
    if let Some(as_of) = cost.pricing_as_of.as_ref() {
        attributes.push(opentelemetry::KeyValue::new(
            format!("nemo_relay.llm.optimization.{label}_pricing_as_of"),
            as_of.clone(),
        ));
    }
}

fn push_optimization_savings_and_status(
    attributes: &mut Vec<opentelemetry::KeyValue>,
    summary: &crate::codec::optimization::LlmOptimizationSummary,
) {
    if let Some(saved) = summary.estimated_cost_saved {
        attributes.push(opentelemetry::KeyValue::new(
            "nemo_relay.llm.optimization.estimated_cost_saved",
            saved,
        ));
        if let Some(currency) = summary.currency.as_ref() {
            attributes.push(opentelemetry::KeyValue::new(
                "nemo_relay.llm.optimization.estimated_cost_saved_currency",
                currency.clone(),
            ));
        }
    }
    if let Some(currency) = summary.currency.as_ref() {
        attributes.push(opentelemetry::KeyValue::new(
            "nemo_relay.llm.optimization.currency",
            currency.clone(),
        ));
    }
    let status = match summary.status {
        crate::codec::optimization::LlmOptimizationSummaryStatus::Complete => "complete",
        crate::codec::optimization::LlmOptimizationSummaryStatus::Partial => "partial",
    };
    attributes.push(opentelemetry::KeyValue::new(
        "nemo_relay.llm.optimization.status",
        status,
    ));
}

fn push_optimization_pricing_provenance(
    attributes: &mut Vec<opentelemetry::KeyValue>,
    summary: &crate::codec::optimization::LlmOptimizationSummary,
) {
    let source = summary
        .baseline_cost
        .as_ref()
        .and_then(|cost| cost.pricing_source.as_ref())
        .or_else(|| {
            summary
                .actual_cost
                .as_ref()
                .and_then(|cost| cost.pricing_source.as_ref())
        });
    if let Some(source) = source {
        attributes.push(opentelemetry::KeyValue::new(
            "nemo_relay.llm.optimization.pricing_source",
            source.clone(),
        ));
    }
    let as_of = summary
        .baseline_cost
        .as_ref()
        .and_then(|cost| cost.pricing_as_of.as_ref())
        .or_else(|| {
            summary
                .actual_cost
                .as_ref()
                .and_then(|cost| cost.pricing_as_of.as_ref())
        });
    if let Some(as_of) = as_of {
        attributes.push(opentelemetry::KeyValue::new(
            "nemo_relay.llm.optimization.pricing_as_of",
            as_of.clone(),
        ));
    }
}

/// Validates OTLP attribute mappings shared by exporter configuration surfaces.
pub fn validate_attribute_mappings(
    mappings: &[OtlpAttributeMapping],
) -> std::result::Result<(), String> {
    let mut aliases = std::collections::HashSet::new();
    for mapping in mappings {
        if is_blank_attribute_mapping_name(&mapping.key) {
            return Err("attribute mapping key must not be blank".to_string());
        }
        if is_blank_attribute_mapping_name(&mapping.alias) {
            return Err("attribute mapping alias must not be blank".to_string());
        }
        if !aliases.insert(mapping.alias.trim()) {
            return Err(format!(
                "attribute mapping alias {:?} is duplicated",
                mapping.alias
            ));
        }
    }
    Ok(())
}

fn is_blank_attribute_mapping_name(value: &str) -> bool {
    value.chars().all(|character| {
        character.is_whitespace()
            || matches!(
                unicode_general_category::get_general_category(character),
                unicode_general_category::GeneralCategory::Control
                    | unicode_general_category::GeneralCategory::Format
            )
    })
}

/// Projects only top-level JSON fields as OTLP attributes.
///
/// Nested objects and arrays remain JSON strings so arbitrary payloads do not
/// create ambiguous dotted attribute paths or unbounded attribute sets.
pub(crate) fn push_top_level_json_attributes(
    attributes: &mut Vec<opentelemetry::KeyValue>,
    prefix: &str,
    value: Option<&crate::json::Json>,
) {
    let Some(value) = value else {
        return;
    };
    match value {
        crate::json::Json::Object(values) => {
            for (field, value) in values {
                push_top_level_json_value(attributes, &format!("{prefix}.{field}"), value);
            }
        }
        value => push_top_level_json_value(attributes, prefix, value),
    }
}

/// Adds canonical session-correlation attributes from event metadata and the
/// active scope-stack instance.
pub(crate) fn push_session_identity_attributes(
    attributes: &mut Vec<opentelemetry::KeyValue>,
    event: &crate::api::event::Event,
) {
    use opentelemetry::KeyValue;

    let metadata = event.metadata();
    if let Some(session_id) = metadata
        .and_then(|value| value.get("session_id"))
        .and_then(crate::json::Json::as_str)
    {
        attributes.push(KeyValue::new("session.id", session_id.to_string()));
    }
    if let Some(user_id) = metadata
        .and_then(|value| value.get("user_id"))
        .and_then(crate::json::Json::as_str)
    {
        attributes.push(KeyValue::new("user.id", user_id.to_string()));
    }
    if let Some(agent_kind) = metadata
        .and_then(|value| value.get("agent_kind"))
        .and_then(crate::json::Json::as_str)
    {
        attributes.push(KeyValue::new(
            "nemo_relay.agent.kind",
            agent_kind.to_string(),
        ));
    }
    if let Ok(stack) = crate::api::runtime::current_scope_stack().read() {
        attributes.push(KeyValue::new(
            "nemo_relay.session.instance_id",
            stack.root_uuid().to_string(),
        ));
    }
}

/// Serializes a value and projects its top-level JSON fields as OTLP attributes.
pub(crate) fn push_serialized_top_level_attributes<T: Serialize + ?Sized>(
    attributes: &mut Vec<opentelemetry::KeyValue>,
    prefix: &str,
    value: Option<&T>,
) {
    let Some(value) = value else {
        return;
    };
    if let Ok(value) = serde_json::to_value(value) {
        push_top_level_json_attributes(attributes, prefix, Some(&value));
    }
}

fn push_top_level_json_value(
    attributes: &mut Vec<opentelemetry::KeyValue>,
    key: &str,
    value: &crate::json::Json,
) {
    use opentelemetry::KeyValue;

    match value {
        crate::json::Json::Null => {}
        crate::json::Json::Bool(value) => attributes.push(KeyValue::new(key.to_string(), *value)),
        crate::json::Json::String(value) => {
            attributes.push(KeyValue::new(key.to_string(), value.clone()))
        }
        crate::json::Json::Number(value) => {
            if let Some(value) = value.as_i64() {
                attributes.push(KeyValue::new(key.to_string(), value));
            } else if let Some(value) = value.as_u64() {
                if let Ok(value) = i64::try_from(value) {
                    attributes.push(KeyValue::new(key.to_string(), value));
                } else {
                    attributes.push(KeyValue::new(key.to_string(), value.to_string()));
                }
            } else if let Some(value) = value.as_f64() {
                attributes.push(KeyValue::new(key.to_string(), value));
            }
        }
        crate::json::Json::Array(_) | crate::json::Json::Object(_) => {
            if let Ok(value) = serde_json::to_string(value) {
                attributes.push(KeyValue::new(key.to_string(), value));
            }
        }
    }
}

pub(crate) fn apply_attribute_mappings(
    attributes: &mut Vec<opentelemetry::KeyValue>,
    mappings: &[OtlpAttributeMapping],
) {
    attributes.extend(attribute_mapping_aliases(attributes, mappings));
}

/// Keeps the start attributes needed to resolve mappings at the end of a span.
///
/// The final span attributes must still take precedence over mapped aliases, so
/// retain both mapped source keys and aliases that were already present at
/// start. The span itself owns all other start attributes and does not need a
/// second copy in the active-span state.
pub(crate) fn attribute_mapping_inputs(
    attributes: &[opentelemetry::KeyValue],
    mappings: &[OtlpAttributeMapping],
) -> Vec<opentelemetry::KeyValue> {
    attributes
        .iter()
        .filter(|attribute| {
            mappings.iter().any(|mapping| {
                attribute.key.as_str() == mapping.key || attribute.key.as_str() == mapping.alias
            })
        })
        .cloned()
        .collect()
}

/// Resolves typed aliases from a complete set of projected attributes.
///
/// Callers that project a span across multiple lifecycle events must pass every
/// real span attribute so projected fields always take precedence over aliases.
pub(crate) fn attribute_mapping_aliases(
    projected_attributes: &[opentelemetry::KeyValue],
    mappings: &[OtlpAttributeMapping],
) -> Vec<opentelemetry::KeyValue> {
    if mappings.is_empty() {
        return Vec::new();
    }
    let existing = projected_attributes
        .iter()
        .map(|attribute| attribute.key.as_str().to_string())
        .collect::<std::collections::HashSet<_>>();
    mappings
        .iter()
        .filter(|mapping| !existing.contains(mapping.alias.as_str()))
        .filter_map(|mapping| {
            projected_attributes
                .iter()
                .rev()
                .find(|attribute| attribute.key.as_str() == mapping.key)
                .map(|attribute| {
                    opentelemetry::KeyValue::new(mapping.alias.clone(), attribute.value.clone())
                })
        })
        .collect()
}

/// Returns whether a mark matches a configured projection exclusion.
///
/// Agent hook adapters may preserve the canonical event name in metadata while
/// using a generic mark name, so both representations are matched.
pub(crate) fn mark_name_is_excluded(
    event: &crate::api::event::Event,
    excluded_names: &[String],
) -> bool {
    excluded_names.iter().any(|name| {
        event.name() == name
            || event
                .metadata()
                .and_then(crate::json::Json::as_object)
                .and_then(|metadata| metadata.get("hook_event_name"))
                .and_then(crate::json::Json::as_str)
                == Some(name.as_str())
    })
}

/// Resolves a configured mark projection for one event.
///
/// Exclusions only affect tool projection; all other modes retain their
/// configured exporter-native behavior.
pub(crate) fn effective_mark_projection(
    event: &crate::api::event::Event,
    projection: MarkProjection,
    excluded_names: &[String],
) -> MarkProjection {
    if projection == MarkProjection::Tool && mark_name_is_excluded(event, excluded_names) {
        MarkProjection::Inherit
    } else {
        projection
    }
}

#[cfg(test)]
#[path = "../../tests/unit/observability/exporter_parity_tests.rs"]
mod exporter_parity_tests;

pub(crate) fn estimate_cost_for_response_or_requested_model(
    event: &crate::api::event::Event,
    response_model: Option<&str>,
    usage: &crate::codec::response::Usage,
) -> Option<crate::codec::response::CostEstimate> {
    estimate_cost_for_response_or_model(
        Some(event.name()),
        event.model_name(),
        response_model,
        usage,
    )
}

pub(crate) fn estimate_cost_for_response_or_model(
    provider: Option<&str>,
    requested_model: Option<&str>,
    response_model: Option<&str>,
    usage: &crate::codec::response::Usage,
) -> Option<crate::codec::response::CostEstimate> {
    // Prefer the provider-echoed model, but fall back to the requested model
    // when pricing does not recognize the echoed model alias.
    if let Some(model_name) = response_model
        && let Some(cost) =
            crate::codec::response::estimate_cost_for_provider(provider, model_name, usage)
    {
        return Some(cost);
    }

    let requested_model = requested_model?;
    if response_model == Some(requested_model) {
        return None;
    }
    crate::codec::response::estimate_cost_for_provider(provider, requested_model, usage)
}

pub(crate) fn merge_usage(
    primary: Option<&crate::codec::response::Usage>,
    secondary: Option<&crate::codec::response::Usage>,
) -> Option<crate::codec::response::Usage> {
    match (primary, secondary) {
        (None, None) => None,
        (None, Some(usage)) | (Some(usage), None) => Some(usage.clone()),
        (Some(primary), Some(secondary)) => Some(crate::codec::response::Usage {
            prompt_tokens: primary.prompt_tokens.or(secondary.prompt_tokens),
            completion_tokens: primary.completion_tokens.or(secondary.completion_tokens),
            total_tokens: primary.total_tokens.or(secondary.total_tokens),
            cache_read_tokens: primary.cache_read_tokens.or(secondary.cache_read_tokens),
            cache_write_tokens: primary.cache_write_tokens.or(secondary.cache_write_tokens),
            cost: primary.cost.clone().or_else(|| secondary.cost.clone()),
        }),
    }
}

pub(crate) fn model_name_for_llm_event(event: &crate::api::event::Event) -> Option<String> {
    if event.category().map(|category| category.as_str()) != Some("llm") {
        return None;
    }
    let manual_response_model =
        manual::model_name_from_manual_llm_output(event.output()).map(ToOwned::to_owned);
    let manual_request_model =
        manual::model_name_from_manual_llm_output(event.input()).map(ToOwned::to_owned);
    event
        .normalized_llm_response()
        .and_then(|response| response.as_ref().model.clone())
        .or(manual_response_model)
        .or_else(|| event.model_name().map(ToOwned::to_owned))
        .or_else(|| {
            event
                .normalized_llm_request()
                .and_then(|request| request.as_ref().model.clone())
        })
        .or(manual_request_model)
}

pub(crate) fn set_span_status_from_event_metadata<S>(span: &mut S, event: &crate::api::event::Event)
where
    S: opentelemetry::trace::Span,
{
    let Some(metadata) = event.metadata() else {
        return;
    };
    let Some(status_code) = metadata
        .get("otel.status_code")
        .and_then(crate::json::Json::as_str)
    else {
        return;
    };

    let status = match status_code {
        "OK" => opentelemetry::trace::Status::Ok,
        "ERROR" => opentelemetry::trace::Status::error(
            metadata
                .get("otel.status_description")
                .and_then(crate::json::Json::as_str)
                .unwrap_or_default()
                .to_string(),
        ),
        "UNSET" => opentelemetry::trace::Status::Unset,
        _ => {
            log::warn!(
                target: "nemo_relay.observability",
                event = "invalid_status_code",
                status_code = "invalid";
                "Unrecognized OpenTelemetry status code; using unset status"
            );
            opentelemetry::trace::Status::Unset
        }
    };
    span.set_status(status);
}

#[cfg(test)]
#[path = "../../tests/unit/observability/attribute_projection_tests.rs"]
mod attribute_projection_tests;

#[cfg(test)]
#[path = "../../tests/unit/observability/mod_tests.rs"]
mod tests;