Skip to main content

nemo_relay/observability/
otel.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! OpenTelemetry subscriber support for NeMo Relay.
5//!
6//! This crate adapts NeMo Relay lifecycle events into the selected `full`,
7//! `gen_ai`, or `openinference` OpenTelemetry trace projection. Scope start and
8//! end events open and close supported spans. Mark behavior is fixed by the
9//! selected projection.
10//!
11//! The public API is intentionally small:
12//!
13//! - [`OpenTelemetryConfig`] configures the OTLP exporter and resource metadata
14//! - [`OpenTelemetrySubscriber`] exposes a NeMo Relay [`EventSubscriberFn`] and
15//!   convenience `register` / `deregister` / `force_flush` / `shutdown` methods
16
17use std::borrow::Cow;
18use std::cell::RefCell;
19use std::collections::{HashMap, HashSet, VecDeque};
20use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
21use std::sync::mpsc;
22use std::sync::{Arc, Mutex};
23use std::thread;
24use std::time::{Duration, SystemTime, UNIX_EPOCH};
25
26use super::{
27    MarkProjection, OpenTelemetryType, OtlpAttributeMapping, apply_attribute_mappings,
28    attribute_mapping_aliases, attribute_mapping_inputs, default_mark_exclude_names,
29    effective_mark_projection, estimate_cost_for_response_or_model,
30    estimate_cost_for_response_or_requested_model, manual, model_name_for_llm_event,
31    push_serialized_top_level_attributes, push_session_identity_attributes,
32    push_top_level_json_attributes, relay_span_id, relay_trace_id, validate_attribute_mappings,
33};
34use crate::api::event::{Event, EventNormalizationExt, ScopeCategory};
35use crate::api::runtime::{EventSubscriberFn, current_scope_stack};
36use crate::api::scope::ScopeType;
37use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber};
38use crate::codec::response::CostEstimate;
39use crate::error::FlowError;
40use chrono::{DateTime, Utc};
41use opentelemetry::trace::{
42    Span as _, SpanContext, SpanId, SpanKind, TraceContextExt, TraceFlags, TraceId, TraceState,
43    Tracer, TracerProvider as _,
44};
45use opentelemetry::{Context, KeyValue};
46use opentelemetry_otlp::{
47    Protocol, SpanExporter as OtlpSpanExporter, WithExportConfig, WithHttpConfig,
48};
49use opentelemetry_sdk::Resource;
50use opentelemetry_sdk::error::{OTelSdkError, OTelSdkResult};
51use opentelemetry_sdk::trace::{
52    BatchSpanProcessor, IdGenerator, RandomIdGenerator, SdkTracer, SdkTracerProvider, Span,
53    SpanData, SpanExporter, SpanProcessor,
54};
55use uuid::Uuid;
56
57use crate::plugin::{
58    OTEL_RUNTIME_DELIVERY_FAILURE_MARKER, RuntimeDiagnostic,
59    record_active_plugin_runtime_diagnostic,
60};
61
62pub(super) const COMPLETED_SPAN_CONTEXT_LIMIT: usize = 4096;
63
64use opentelemetry_otlp::WithTonicConfig;
65use tonic::metadata::{MetadataKey, MetadataMap, MetadataValue};
66
67thread_local! {
68    static PENDING_RELAY_IDS: RefCell<Option<(TraceId, SpanId)>> = const { RefCell::new(None) };
69}
70
71#[derive(Clone, Debug, Default)]
72pub(crate) struct RelayIdGenerator;
73
74impl IdGenerator for RelayIdGenerator {
75    fn new_trace_id(&self) -> TraceId {
76        PENDING_RELAY_IDS
77            .with(|ids| ids.borrow().map(|(trace_id, _)| trace_id))
78            .unwrap_or_else(|| RandomIdGenerator::default().new_trace_id())
79    }
80
81    fn new_span_id(&self) -> SpanId {
82        PENDING_RELAY_IDS
83            .with(|ids| ids.borrow().map(|(_, span_id)| span_id))
84            .unwrap_or_else(|| RandomIdGenerator::default().new_span_id())
85    }
86}
87
88fn with_relay_ids<T>(uuid: Uuid, build: impl FnOnce() -> T) -> T {
89    struct ResetPendingIds;
90
91    impl Drop for ResetPendingIds {
92        fn drop(&mut self) {
93            PENDING_RELAY_IDS.with(|ids| {
94                ids.replace(None);
95            });
96        }
97    }
98
99    PENDING_RELAY_IDS.with(|ids| {
100        ids.replace(Some((
101            super::relay_trace_id(uuid),
102            super::relay_span_id(uuid),
103        )));
104    });
105    let _reset = ResetPendingIds;
106    build()
107}
108
109/// Result type for the OpenTelemetry subscriber crate.
110pub type Result<T> = std::result::Result<T, OpenTelemetryError>;
111
112/// Errors produced while configuring or operating the OpenTelemetry subscriber.
113#[derive(Debug, thiserror::Error)]
114pub enum OpenTelemetryError {
115    /// Failed to parse a configured gRPC metadata header.
116    #[error("invalid OTLP gRPC header {key:?}: {message}")]
117    InvalidGrpcHeader {
118        /// Header name that failed to parse.
119        key: String,
120        /// Parser failure message.
121        message: String,
122    },
123    /// A configured OTLP header name or value is invalid.
124    #[error("invalid OTLP header {key:?}: {message}")]
125    InvalidHeader {
126        /// Header name that failed validation.
127        key: String,
128        /// Validation failure message.
129        message: String,
130    },
131    /// Process-global OTLP header variables cannot be isolated per endpoint.
132    #[error(
133        "{variable} is not supported because process-global OTLP headers can leak across endpoints; use the endpoint headers or header_env configuration"
134    )]
135    GlobalHeaderEnvironmentUnsupported {
136        /// Environment variable that was set.
137        variable: &'static str,
138    },
139    /// Failed to build the OTLP exporter.
140    #[error("failed to build the OTLP exporter: {0}")]
141    ExporterBuild(String),
142    /// The underlying tracer provider returned an error.
143    #[error("OpenTelemetry tracer provider error: {0}")]
144    Provider(String),
145    /// Attribute mapping configuration was invalid.
146    #[error("invalid attribute mappings: {0}")]
147    InvalidAttributeMappings(String),
148    /// Registration errors from the core runtime.
149    #[error(transparent)]
150    Core(#[from] FlowError),
151}
152
153/// Supported OTLP trace transports.
154#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
155pub enum OtlpTransport {
156    /// OTLP/HTTP protobuf, typically `http://host:4318/v1/traces`.
157    #[default]
158    HttpBinary,
159    /// OTLP/gRPC, typically `http://host:4317`.
160    Grpc,
161}
162
163/// Completes a bare OTLP/HTTP base URL with the standard trace signal path.
164#[doc(hidden)]
165pub fn resolve_http_trace_endpoint(endpoint: &str) -> Cow<'_, str> {
166    let Ok(mut parsed) = reqwest::Url::parse(endpoint) else {
167        return Cow::Borrowed(endpoint);
168    };
169
170    let has_explicit_root_path = endpoint
171        .split(['?', '#'])
172        .next()
173        .is_some_and(|url| url.ends_with('/'));
174    if !matches!(parsed.scheme(), "http" | "https")
175        || parsed.path() != "/"
176        || has_explicit_root_path
177    {
178        return Cow::Borrowed(endpoint);
179    }
180    parsed.set_path("/v1/traces");
181    Cow::Owned(parsed.into())
182}
183
184/// Configuration for the OpenTelemetry subscriber.
185#[derive(Debug, Clone)]
186pub struct OpenTelemetryConfig {
187    otel_type: OpenTelemetryType,
188    endpoint: String,
189    headers: HashMap<String, String>,
190    resource_attributes: HashMap<String, String>,
191    service_name: String,
192    service_namespace: Option<String>,
193    service_version: Option<String>,
194    instrumentation_scope: String,
195    mark_projection: MarkProjection,
196    mark_exclude_names: Vec<String>,
197    attribute_mappings: Vec<OtlpAttributeMapping>,
198    timeout: Duration,
199    transport: OtlpTransport,
200}
201
202impl OpenTelemetryConfig {
203    fn default_values() -> Self {
204        Self {
205            otel_type: OpenTelemetryType::Full,
206            endpoint: String::new(),
207            headers: HashMap::new(),
208            resource_attributes: HashMap::new(),
209            service_name: "unknown_service".to_string(),
210            service_namespace: None,
211            service_version: None,
212            instrumentation_scope: "opentelemetry".to_string(),
213            mark_projection: MarkProjection::default(),
214            mark_exclude_names: default_mark_exclude_names(),
215            attribute_mappings: Vec::new(),
216            timeout: Duration::from_secs(3),
217            transport: OtlpTransport::HttpBinary,
218        }
219    }
220
221    /// Creates a typed OpenTelemetry exporter for a required OTLP endpoint.
222    pub fn new(otel_type: OpenTelemetryType, endpoint: impl Into<String>) -> Self {
223        Self {
224            otel_type,
225            endpoint: endpoint.into(),
226            ..Self::default_values()
227        }
228    }
229
230    /// Creates an HTTP OTLP config for the given service name.
231    #[cfg(test)]
232    pub(crate) fn http_binary(service_name: impl Into<String>) -> Self {
233        Self {
234            service_name: service_name.into(),
235            transport: OtlpTransport::HttpBinary,
236            ..Self::default_values()
237        }
238    }
239
240    /// Creates a gRPC OTLP config for the given service name.
241    #[cfg(test)]
242    pub(crate) fn grpc(service_name: impl Into<String>) -> Self {
243        Self {
244            service_name: service_name.into(),
245            transport: OtlpTransport::Grpc,
246            ..Self::default_values()
247        }
248    }
249
250    /// Overrides the OTLP endpoint. If unset, exporter defaults and OTEL_* env vars apply.
251    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
252        self.endpoint = endpoint.into();
253        self
254    }
255
256    /// Selects the OTLP transport.
257    pub fn with_transport(mut self, transport: OtlpTransport) -> Self {
258        self.transport = transport;
259        self
260    }
261
262    /// Sets the `service.name` resource attribute.
263    pub fn with_service_name(mut self, service_name: impl Into<String>) -> Self {
264        self.service_name = service_name.into();
265        self
266    }
267
268    /// Adds a header/metadata entry for the exporter.
269    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
270        self.headers.insert(key.into(), value.into());
271        self
272    }
273
274    #[cfg(test)]
275    pub(crate) fn header(&self, key: &str) -> Option<&str> {
276        self.headers.get(key).map(String::as_str)
277    }
278
279    /// Adds a resource attribute as a string key/value pair.
280    pub fn with_resource_attribute(
281        mut self,
282        key: impl Into<String>,
283        value: impl Into<String>,
284    ) -> Self {
285        self.resource_attributes.insert(key.into(), value.into());
286        self
287    }
288
289    /// Sets the OTLP request timeout.
290    pub fn with_timeout(mut self, timeout: Duration) -> Self {
291        self.timeout = timeout;
292        self
293    }
294
295    /// Sets the service namespace resource attribute.
296    pub fn with_service_namespace(mut self, namespace: impl Into<String>) -> Self {
297        self.service_namespace = Some(namespace.into());
298        self
299    }
300
301    /// Sets the service version resource attribute.
302    pub fn with_service_version(mut self, version: impl Into<String>) -> Self {
303        self.service_version = Some(version.into());
304        self
305    }
306
307    /// Sets the instrumentation scope name used for emitted spans.
308    pub fn with_instrumentation_scope(mut self, scope: impl Into<String>) -> Self {
309        self.instrumentation_scope = scope.into();
310        self
311    }
312
313    /// Selects how point-in-time marks are represented in exported traces.
314    pub fn with_mark_projection(mut self, mark_projection: MarkProjection) -> Self {
315        self.mark_projection = mark_projection;
316        self
317    }
318
319    /// Excludes named marks from tool projection.
320    pub fn with_mark_exclude_names<I, S>(mut self, names: I) -> Self
321    where
322        I: IntoIterator<Item = S>,
323        S: Into<String>,
324    {
325        self.mark_exclude_names = names.into_iter().map(Into::into).collect();
326        self
327    }
328
329    /// Adds a projected OpenTelemetry attribute alias.
330    pub fn with_attribute_mapping(
331        mut self,
332        key: impl Into<String>,
333        alias: impl Into<String>,
334    ) -> Self {
335        self.attribute_mappings
336            .push(OtlpAttributeMapping::new(key, alias));
337        self
338    }
339
340    /// Replaces projected OpenTelemetry attribute aliases.
341    pub fn with_attribute_mappings<I>(mut self, mappings: I) -> Self
342    where
343        I: IntoIterator<Item = OtlpAttributeMapping>,
344    {
345        self.attribute_mappings = mappings.into_iter().collect();
346        self
347    }
348}
349
350#[cfg(test)]
351impl Default for OpenTelemetryConfig {
352    fn default() -> Self {
353        Self::default_values()
354    }
355}
356
357/// OpenTelemetry-backed NeMo Relay subscriber.
358#[derive(Clone)]
359pub struct OpenTelemetrySubscriber {
360    inner: Arc<Inner>,
361}
362
363/// Options for constructing an OpenTelemetry subscriber from an existing tracer provider.
364#[derive(Debug, Clone)]
365pub struct OpenTelemetrySubscriberOptions {
366    /// How mark events are projected into the trace.
367    pub mark_projection: MarkProjection,
368    /// Mark names excluded from tool projection.
369    pub mark_exclude_names: Vec<String>,
370    /// Typed OTLP attributes copied to alias keys.
371    pub attribute_mappings: Vec<OtlpAttributeMapping>,
372}
373
374impl Default for OpenTelemetrySubscriberOptions {
375    fn default() -> Self {
376        Self {
377            mark_projection: MarkProjection::default(),
378            mark_exclude_names: default_mark_exclude_names(),
379            attribute_mappings: Vec::new(),
380        }
381    }
382}
383
384struct Inner {
385    // Keep `processor` before `_runtime`: the processor owns the
386    // `SdkTracerProvider` and must be dropped before `ExporterRuntime` joins
387    // and tears down its Tokio runtime. Do not reorder these fields.
388    processor: Arc<Mutex<OtelEventProcessor>>,
389    subscriber: EventSubscriberFn,
390    _runtime: Option<ExporterRuntime>,
391}
392
393struct ExporterRuntime {
394    stop: Option<mpsc::Sender<()>>,
395    thread: Option<thread::JoinHandle<()>>,
396}
397
398impl Drop for ExporterRuntime {
399    fn drop(&mut self) {
400        self.stop.take();
401        if let Some(thread) = self.thread.take() {
402            let _ = thread.join();
403        }
404    }
405}
406
407impl OpenTelemetrySubscriber {
408    /// Builds a subscriber backed by a new OTLP tracer provider.
409    pub fn new(config: OpenTelemetryConfig) -> Result<Self> {
410        Self::new_with_runtime_diagnostics(config, None)
411    }
412
413    pub(crate) fn new_for_plugin(
414        config: OpenTelemetryConfig,
415        endpoint_index: usize,
416    ) -> Result<Self> {
417        Self::new_with_runtime_diagnostics(
418            config,
419            Some(format!(
420                "opentelemetry.endpoints[{endpoint_index}].endpoint"
421            )),
422        )
423    }
424
425    fn new_with_runtime_diagnostics(
426        config: OpenTelemetryConfig,
427        diagnostic_field: Option<String>,
428    ) -> Result<Self> {
429        if config.endpoint.trim().is_empty() {
430            return Err(OpenTelemetryError::ExporterBuild(
431                "endpoint must be a nonblank string".to_string(),
432            ));
433        }
434        validate_attribute_mappings(&config.attribute_mappings)
435            .map_err(OpenTelemetryError::InvalidAttributeMappings)?;
436        reject_global_header_environment()?;
437        validate_headers(&config.headers)?;
438        let (provider, runtime) = build_owned_tracer_provider(config.clone(), diagnostic_field)?;
439        Ok(Self::from_tracer_provider_with_scope_and_type(
440            provider,
441            config.instrumentation_scope,
442            config.otel_type,
443            config.mark_projection,
444            config.mark_exclude_names,
445            config.attribute_mappings,
446            Some(runtime),
447        ))
448    }
449
450    /// Builds a subscriber from an already-configured tracer provider.
451    pub fn from_tracer_provider(
452        provider: SdkTracerProvider,
453        instrumentation_scope: impl Into<String>,
454    ) -> Self {
455        Self::from_tracer_provider_with_type(
456            provider,
457            instrumentation_scope,
458            OpenTelemetryType::Full,
459        )
460    }
461
462    /// Builds a typed subscriber from an already-configured tracer provider.
463    pub fn from_tracer_provider_with_type(
464        provider: SdkTracerProvider,
465        instrumentation_scope: impl Into<String>,
466        otel_type: OpenTelemetryType,
467    ) -> Self {
468        let instrumentation_scope = instrumentation_scope.into();
469        Self::from_tracer_provider_with_scope_and_type(
470            provider,
471            instrumentation_scope,
472            otel_type,
473            MarkProjection::default(),
474            default_mark_exclude_names(),
475            Vec::new(),
476            None,
477        )
478    }
479
480    /// Builds a subscriber from a tracer provider with typed attribute copies.
481    pub fn from_tracer_provider_with_attribute_mappings<I>(
482        provider: SdkTracerProvider,
483        instrumentation_scope: impl Into<String>,
484        attribute_mappings: I,
485    ) -> Result<Self>
486    where
487        I: IntoIterator<Item = OtlpAttributeMapping>,
488    {
489        let attribute_mappings = attribute_mappings.into_iter().collect::<Vec<_>>();
490        Self::from_tracer_provider_with_options(
491            provider,
492            instrumentation_scope,
493            OpenTelemetrySubscriberOptions {
494                attribute_mappings,
495                ..Default::default()
496            },
497        )
498    }
499
500    /// Builds a subscriber from a tracer provider with composable projection options.
501    pub fn from_tracer_provider_with_options(
502        provider: SdkTracerProvider,
503        instrumentation_scope: impl Into<String>,
504        options: OpenTelemetrySubscriberOptions,
505    ) -> Result<Self> {
506        validate_attribute_mappings(&options.attribute_mappings)
507            .map_err(OpenTelemetryError::InvalidAttributeMappings)?;
508        Ok(Self::from_tracer_provider_with_scope_and_type(
509            provider,
510            instrumentation_scope.into(),
511            OpenTelemetryType::Full,
512            options.mark_projection,
513            options.mark_exclude_names,
514            options.attribute_mappings,
515            None,
516        ))
517    }
518
519    /// Builds a typed subscriber from a tracer provider with projection options.
520    pub fn from_tracer_provider_with_type_and_options(
521        provider: SdkTracerProvider,
522        instrumentation_scope: impl Into<String>,
523        otel_type: OpenTelemetryType,
524        options: OpenTelemetrySubscriberOptions,
525    ) -> Result<Self> {
526        validate_attribute_mappings(&options.attribute_mappings)
527            .map_err(OpenTelemetryError::InvalidAttributeMappings)?;
528        Ok(Self::from_tracer_provider_with_scope_and_type(
529            provider,
530            instrumentation_scope.into(),
531            otel_type,
532            options.mark_projection,
533            options.mark_exclude_names,
534            options.attribute_mappings,
535            None,
536        ))
537    }
538
539    fn from_tracer_provider_with_scope_and_type(
540        provider: SdkTracerProvider,
541        instrumentation_scope: String,
542        otel_type: OpenTelemetryType,
543        mark_projection: MarkProjection,
544        mark_exclude_names: Vec<String>,
545        attribute_mappings: Vec<OtlpAttributeMapping>,
546        runtime: Option<ExporterRuntime>,
547    ) -> Self {
548        let processor = Arc::new(Mutex::new(
549            OtelEventProcessor::new_with_mark_projection_and_exclusions_and_mappings(
550                provider,
551                instrumentation_scope,
552                otel_type,
553                mark_projection,
554                mark_exclude_names,
555                attribute_mappings,
556            ),
557        ));
558        let processor_for_callback = Arc::clone(&processor);
559        let subscriber: EventSubscriberFn = Arc::new(move |event: &Event| {
560            let Ok(mut guard) = processor_for_callback.lock() else {
561                // Observability should not take down the host process if the
562                // subscriber state was previously poisoned.
563                return;
564            };
565            guard.process(event);
566        });
567
568        Self {
569            inner: Arc::new(Inner {
570                processor,
571                subscriber,
572                _runtime: runtime,
573            }),
574        }
575    }
576
577    /// Returns the raw NeMo Relay subscriber callback for custom registration flows.
578    pub fn subscriber(&self) -> EventSubscriberFn {
579        Arc::clone(&self.inner.subscriber)
580    }
581
582    /// Registers this subscriber globally with the NeMo Relay runtime.
583    pub fn register(&self, name: &str) -> Result<()> {
584        register_subscriber(name, self.subscriber())?;
585        log::info!(
586            target: "nemo_relay.observability",
587            event = "exporter_registered",
588            exporter = "opentelemetry",
589            subscriber = name;
590            "OpenTelemetry exporter registered"
591        );
592        Ok(())
593    }
594
595    /// Deregisters a previously-registered global subscriber by name.
596    pub fn deregister(&self, name: &str) -> Result<bool> {
597        let removed = deregister_subscriber(name)?;
598        if removed {
599            log::info!(
600                target: "nemo_relay.observability",
601                event = "subscriber_deregistered",
602                subscriber = name;
603                "Observability subscriber deregistered"
604            );
605        }
606        Ok(removed)
607    }
608
609    /// Flushes finished spans through the underlying tracer provider.
610    pub fn force_flush(&self) -> Result<()> {
611        flush_subscribers()?;
612        let guard = self.inner.processor.lock().map_err(|_| {
613            OpenTelemetryError::Provider("the subscriber state lock was poisoned".to_string())
614        })?;
615        guard.force_flush()
616    }
617
618    /// Shuts down the underlying tracer provider.
619    ///
620    /// Call `deregister(...)` first if the subscriber is still registered with NeMo Relay.
621    pub fn shutdown(&self) -> Result<()> {
622        let barrier_error = flush_subscribers().err().map(OpenTelemetryError::Core);
623        let provider_result = self.shutdown_provider();
624        if let Some(error) = barrier_error {
625            return Err(error);
626        }
627        provider_result
628    }
629
630    pub(crate) fn shutdown_provider(&self) -> Result<()> {
631        let guard = self.inner.processor.lock().map_err(|_| {
632            OpenTelemetryError::Provider("the subscriber state lock was poisoned".to_string())
633        })?;
634        let result = guard.shutdown();
635        if result.is_ok() {
636            log::info!(
637                target: "nemo_relay.observability",
638                event = "exporter_shutdown",
639                exporter = "opentelemetry";
640                "OpenTelemetry exporter shut down"
641            );
642        }
643        result
644    }
645}
646
647fn build_owned_tracer_provider(
648    config: OpenTelemetryConfig,
649    diagnostic_field: Option<String>,
650) -> Result<(SdkTracerProvider, ExporterRuntime)> {
651    let (result_sender, result_receiver) = mpsc::sync_channel(1);
652    let (stop_sender, stop_receiver) = mpsc::channel();
653    let runtime_thread = thread::Builder::new()
654        .name("nemo-relay-otlp".to_string())
655        .spawn(move || {
656            let runtime = match tokio::runtime::Builder::new_multi_thread()
657                .worker_threads(1)
658                .enable_all()
659                .build()
660            {
661                Ok(runtime) => runtime,
662                Err(error) => {
663                    let _ = result_sender
664                        .send(Err(OpenTelemetryError::ExporterBuild(error.to_string())));
665                    return;
666                }
667            };
668            let provider = {
669                let _guard = runtime.enter();
670                build_tracer_provider(&config, diagnostic_field)
671            };
672            let keep_runtime_alive = provider.is_ok();
673            let _ = result_sender.send(provider);
674            if keep_runtime_alive {
675                let _ = stop_receiver.recv();
676            }
677        })
678        .map_err(|error| OpenTelemetryError::ExporterBuild(error.to_string()))?;
679    let provider = result_receiver.recv().map_err(|error| {
680        OpenTelemetryError::ExporterBuild(format!("exporter runtime stopped unexpectedly: {error}"))
681    })??;
682    Ok((
683        provider,
684        ExporterRuntime {
685            stop: Some(stop_sender),
686            thread: Some(runtime_thread),
687        },
688    ))
689}
690
691fn reject_global_header_environment() -> Result<()> {
692    for variable in [
693        "OTEL_EXPORTER_OTLP_HEADERS",
694        "OTEL_EXPORTER_OTLP_TRACES_HEADERS",
695    ] {
696        if std::env::var_os(variable).is_some_and(|value| !value.is_empty()) {
697            return Err(OpenTelemetryError::GlobalHeaderEnvironmentUnsupported { variable });
698        }
699    }
700    Ok(())
701}
702
703pub(crate) fn validate_headers(headers: &HashMap<String, String>) -> Result<()> {
704    let mut normalized = HashSet::new();
705    for (key, value) in headers {
706        let normalized_key = key.to_ascii_lowercase();
707        if !normalized.insert(normalized_key) {
708            return Err(OpenTelemetryError::InvalidHeader {
709                key: key.clone(),
710                message: "header names must be unique ignoring ASCII case".to_string(),
711            });
712        }
713        reqwest::header::HeaderName::from_bytes(key.as_bytes()).map_err(|error| {
714            OpenTelemetryError::InvalidHeader {
715                key: key.clone(),
716                message: error.to_string(),
717            }
718        })?;
719        reqwest::header::HeaderValue::from_str(value).map_err(|error| {
720            OpenTelemetryError::InvalidHeader {
721                key: key.clone(),
722                message: error.to_string(),
723            }
724        })?;
725    }
726    Ok(())
727}
728
729fn build_tracer_provider(
730    config: &OpenTelemetryConfig,
731    diagnostic_field: Option<String>,
732) -> Result<SdkTracerProvider> {
733    let exporter = match config.transport {
734        OtlpTransport::HttpBinary => {
735            let mut builder = OtlpSpanExporter::builder()
736                .with_http()
737                .with_protocol(Protocol::HttpBinary)
738                .with_timeout(config.timeout);
739            builder =
740                builder.with_endpoint(resolve_http_trace_endpoint(&config.endpoint).into_owned());
741            if !config.headers.is_empty() {
742                builder = builder.with_headers(config.headers.clone());
743            }
744            builder
745                .build()
746                .map_err(|e| OpenTelemetryError::ExporterBuild(e.to_string()))?
747        }
748        OtlpTransport::Grpc => {
749            let mut builder = OtlpSpanExporter::builder()
750                .with_tonic()
751                .with_protocol(Protocol::Grpc)
752                .with_timeout(config.timeout);
753            builder = builder.with_endpoint(config.endpoint.clone());
754            if !config.headers.is_empty() {
755                builder = builder.with_metadata(build_grpc_metadata(&config.headers)?);
756            }
757            builder
758                .build()
759                .map_err(|e| OpenTelemetryError::ExporterBuild(e.to_string()))?
760        }
761    };
762
763    let mut resource_attributes = vec![KeyValue::new("service.name", config.service_name.clone())];
764    if let Some(service_namespace) = &config.service_namespace {
765        resource_attributes.push(KeyValue::new(
766            "service.namespace",
767            service_namespace.clone(),
768        ));
769    }
770    if let Some(service_version) = &config.service_version {
771        resource_attributes.push(KeyValue::new("service.version", service_version.clone()));
772    }
773    for (key, value) in &config.resource_attributes {
774        resource_attributes.push(KeyValue::new(key.clone(), value.clone()));
775    }
776
777    // Disable per-span attribute caps. Consumers may emit large attribute
778    // sets on long-running spans; the OTel SDK default (128) silently drops
779    // attributes added last in the span's lifecycle.
780    let builder = SdkTracerProvider::builder()
781        .with_resource(
782            Resource::builder_empty()
783                .with_attributes(resource_attributes)
784                .build(),
785        )
786        .with_id_generator(RelayIdGenerator)
787        .with_max_attributes_per_span(u32::MAX)
788        .with_max_attributes_per_event(u32::MAX);
789
790    let processor =
791        DiagnosticBatchSpanProcessor::new(exporter, config.endpoint.clone(), diagnostic_field);
792    Ok(builder.with_span_processor(processor).build())
793}
794
795#[derive(Debug)]
796struct CountingSpanExporter<E> {
797    inner: E,
798    accepted_spans: Arc<AtomicU64>,
799}
800
801impl<E: SpanExporter> SpanExporter for CountingSpanExporter<E> {
802    async fn export(&self, batch: Vec<SpanData>) -> OTelSdkResult {
803        self.accepted_spans
804            .fetch_add(batch.len() as u64, Ordering::Relaxed);
805        self.inner.export(batch).await
806    }
807
808    fn shutdown_with_timeout(&self, timeout: Duration) -> OTelSdkResult {
809        self.inner.shutdown_with_timeout(timeout)
810    }
811
812    fn force_flush(&self) -> OTelSdkResult {
813        self.inner.force_flush()
814    }
815
816    fn set_resource(&mut self, resource: &Resource) {
817        self.inner.set_resource(resource);
818    }
819}
820
821#[derive(Debug)]
822struct DiagnosticBatchSpanProcessor {
823    inner: BatchSpanProcessor,
824    completed_spans: AtomicU64,
825    accepted_spans: Arc<AtomicU64>,
826    endpoint: String,
827    diagnostic_field: Option<String>,
828    diagnostic_reported: AtomicBool,
829}
830
831impl DiagnosticBatchSpanProcessor {
832    fn new<E: SpanExporter + 'static>(
833        exporter: E,
834        endpoint: String,
835        diagnostic_field: Option<String>,
836    ) -> Self {
837        Self::new_with_batch_config(
838            exporter,
839            endpoint,
840            diagnostic_field,
841            opentelemetry_sdk::trace::BatchConfig::default(),
842        )
843    }
844
845    fn new_with_batch_config<E: SpanExporter + 'static>(
846        exporter: E,
847        endpoint: String,
848        diagnostic_field: Option<String>,
849        batch_config: opentelemetry_sdk::trace::BatchConfig,
850    ) -> Self {
851        let accepted_spans = Arc::new(AtomicU64::new(0));
852        let exporter = CountingSpanExporter {
853            inner: exporter,
854            accepted_spans: Arc::clone(&accepted_spans),
855        };
856        Self {
857            inner: BatchSpanProcessor::builder(exporter)
858                .with_batch_config(batch_config)
859                .build(),
860            completed_spans: AtomicU64::new(0),
861            accepted_spans,
862            endpoint,
863            diagnostic_field,
864            diagnostic_reported: AtomicBool::new(false),
865        }
866    }
867
868    fn record_dropped_spans(&self) -> u64 {
869        let dropped = self
870            .completed_spans
871            .load(Ordering::Relaxed)
872            .saturating_sub(self.accepted_spans.load(Ordering::Relaxed));
873        if dropped == 0
874            || self.diagnostic_field.is_none()
875            || self.diagnostic_reported.swap(true, Ordering::Relaxed)
876        {
877            return dropped;
878        }
879        record_active_plugin_runtime_diagnostic(RuntimeDiagnostic {
880            code: "otel.spans_dropped".to_string(),
881            component: "observability".to_string(),
882            field: self.diagnostic_field.clone(),
883            message: format!(
884                "OpenTelemetry dropped {dropped} spans before export to endpoint {} because the batch queue was full",
885                self.endpoint
886            ),
887            session_id: None,
888            count: dropped,
889        });
890        dropped
891    }
892}
893
894impl SpanProcessor for DiagnosticBatchSpanProcessor {
895    fn on_start(&self, span: &mut Span, cx: &Context) {
896        self.inner.on_start(span, cx);
897    }
898
899    fn on_end(&self, span: SpanData) {
900        self.completed_spans.fetch_add(1, Ordering::Relaxed);
901        self.inner.on_end(span);
902    }
903
904    fn force_flush(&self) -> OTelSdkResult {
905        self.inner.force_flush()
906    }
907
908    fn shutdown_with_timeout(&self, timeout: Duration) -> OTelSdkResult {
909        let result = self.inner.shutdown_with_timeout(timeout);
910        if result.is_ok() {
911            let dropped = self.record_dropped_spans();
912            if dropped > 0 && self.diagnostic_field.is_some() {
913                return Err(OTelSdkError::InternalFailure(format!(
914                    "{OTEL_RUNTIME_DELIVERY_FAILURE_MARKER}: otel.spans_dropped ({dropped})"
915                )));
916            }
917        }
918        result
919    }
920
921    fn set_resource(&mut self, resource: &Resource) {
922        self.inner.set_resource(resource);
923    }
924}
925
926fn build_grpc_metadata(headers: &HashMap<String, String>) -> Result<MetadataMap> {
927    let mut metadata = MetadataMap::new();
928    for (key, value) in headers {
929        let metadata_key = MetadataKey::from_bytes(key.as_bytes()).map_err(|e| {
930            OpenTelemetryError::InvalidGrpcHeader {
931                key: key.clone(),
932                message: e.to_string(),
933            }
934        })?;
935        let metadata_value = MetadataValue::try_from(value.as_str()).map_err(|e| {
936            OpenTelemetryError::InvalidGrpcHeader {
937                key: key.clone(),
938                message: e.to_string(),
939            }
940        })?;
941        metadata.insert(metadata_key, metadata_value);
942    }
943    Ok(metadata)
944}
945
946pub(super) struct ActiveSpan {
947    span: Span,
948    span_context: SpanContext,
949    start_model_name: Option<String>,
950    projected_attributes: Vec<KeyValue>,
951    descendant_error_type: Option<String>,
952    descendant_exception_type: Option<String>,
953}
954
955pub(super) struct OtelEventProcessor {
956    pub(super) active_spans: HashMap<Uuid, ActiveSpan>,
957    pub(super) completed_span_contexts: HashMap<Uuid, SpanContext>,
958    pub(super) completed_span_order: VecDeque<Uuid>,
959    provider: SdkTracerProvider,
960    tracer: SdkTracer,
961    otel_type: OpenTelemetryType,
962    mark_projection: MarkProjection,
963    mark_exclude_names: Vec<String>,
964    attribute_mappings: Vec<OtlpAttributeMapping>,
965}
966
967impl OtelEventProcessor {
968    #[cfg(test)]
969    fn new(provider: SdkTracerProvider, instrumentation_scope: String) -> Self {
970        Self::new_with_mark_projection(provider, instrumentation_scope, MarkProjection::default())
971    }
972
973    #[cfg(test)]
974    pub(super) fn new_openinference(
975        provider: SdkTracerProvider,
976        instrumentation_scope: String,
977    ) -> Self {
978        Self::new_with_mark_projection_and_exclusions_and_mappings(
979            provider,
980            instrumentation_scope,
981            OpenTelemetryType::OpenInference,
982            MarkProjection::default(),
983            default_mark_exclude_names(),
984            Vec::new(),
985        )
986    }
987
988    #[cfg(test)]
989    pub(super) fn new_openinference_with_mark_projection(
990        provider: SdkTracerProvider,
991        instrumentation_scope: String,
992        mark_projection: MarkProjection,
993    ) -> Self {
994        Self::new_with_mark_projection_and_exclusions_and_mappings(
995            provider,
996            instrumentation_scope,
997            OpenTelemetryType::OpenInference,
998            mark_projection,
999            default_mark_exclude_names(),
1000            Vec::new(),
1001        )
1002    }
1003
1004    #[cfg(test)]
1005    pub(super) fn new_openinference_with_mark_projection_and_exclusions(
1006        provider: SdkTracerProvider,
1007        instrumentation_scope: String,
1008        mark_projection: MarkProjection,
1009        mark_exclude_names: Vec<String>,
1010    ) -> Self {
1011        Self::new_with_mark_projection_and_exclusions_and_mappings(
1012            provider,
1013            instrumentation_scope,
1014            OpenTelemetryType::OpenInference,
1015            mark_projection,
1016            mark_exclude_names,
1017            Vec::new(),
1018        )
1019    }
1020
1021    #[cfg(test)]
1022    fn new_with_mark_projection(
1023        provider: SdkTracerProvider,
1024        instrumentation_scope: String,
1025        mark_projection: MarkProjection,
1026    ) -> Self {
1027        Self::new_with_mark_projection_and_exclusions(
1028            provider,
1029            instrumentation_scope,
1030            mark_projection,
1031            default_mark_exclude_names(),
1032        )
1033    }
1034
1035    #[cfg(test)]
1036    fn new_with_mark_projection_and_exclusions(
1037        provider: SdkTracerProvider,
1038        instrumentation_scope: String,
1039        mark_projection: MarkProjection,
1040        mark_exclude_names: Vec<String>,
1041    ) -> Self {
1042        Self::new_with_mark_projection_and_exclusions_and_mappings(
1043            provider,
1044            instrumentation_scope,
1045            OpenTelemetryType::Full,
1046            mark_projection,
1047            mark_exclude_names,
1048            Vec::new(),
1049        )
1050    }
1051
1052    fn new_with_mark_projection_and_exclusions_and_mappings(
1053        provider: SdkTracerProvider,
1054        instrumentation_scope: String,
1055        otel_type: OpenTelemetryType,
1056        mark_projection: MarkProjection,
1057        mark_exclude_names: Vec<String>,
1058        attribute_mappings: Vec<OtlpAttributeMapping>,
1059    ) -> Self {
1060        let tracer = provider.tracer(instrumentation_scope);
1061        Self {
1062            active_spans: HashMap::new(),
1063            completed_span_contexts: HashMap::new(),
1064            completed_span_order: VecDeque::new(),
1065            provider,
1066            tracer,
1067            otel_type,
1068            mark_projection,
1069            mark_exclude_names,
1070            attribute_mappings,
1071        }
1072    }
1073
1074    pub(super) fn process(&mut self, event: &Event) {
1075        match event.scope_category() {
1076            Some(ScopeCategory::Start) => self.process_start(event),
1077            Some(ScopeCategory::End) => self.process_end(event),
1078            None => self.process_mark(event),
1079        }
1080    }
1081
1082    pub(super) fn force_flush(&self) -> Result<()> {
1083        self.provider
1084            .force_flush()
1085            .map_err(|e| OpenTelemetryError::Provider(e.to_string()))
1086    }
1087
1088    fn shutdown(&self) -> Result<()> {
1089        self.provider
1090            .shutdown()
1091            .map_err(|e| OpenTelemetryError::Provider(e.to_string()))
1092    }
1093
1094    fn process_start(&mut self, event: &Event) {
1095        self.remove_completed_span_context(event.uuid());
1096        let parent_context = self.parent_context(event);
1097        let is_trace_root = !parent_context.span().span_context().is_valid();
1098        let start_model_name = model_name_for_llm_event(event);
1099        let span_name = match self.otel_type {
1100            OpenTelemetryType::Full => span_name(event),
1101            OpenTelemetryType::GenAi => super::otel_genai::span_name(event),
1102            OpenTelemetryType::OpenInference => super::openinference::span_name(event),
1103        };
1104        let span_kind = match self.otel_type {
1105            OpenTelemetryType::Full => span_kind(event),
1106            OpenTelemetryType::GenAi => super::otel_genai::span_kind(event),
1107            OpenTelemetryType::OpenInference => super::openinference::span_kind(event),
1108        };
1109        let mut span = with_relay_ids(event.uuid(), || {
1110            self.tracer
1111                .span_builder(span_name)
1112                .with_kind(span_kind)
1113                .with_start_time(to_system_time(*event.timestamp()))
1114                .start_with_context(&self.tracer, &parent_context)
1115        });
1116        let mut attributes = match self.otel_type {
1117            OpenTelemetryType::Full => start_attributes(event),
1118            OpenTelemetryType::GenAi => super::otel_genai::start_attributes(event),
1119            OpenTelemetryType::OpenInference => super::openinference::start_attributes(event),
1120        };
1121        if self.otel_type == OpenTelemetryType::Full && start_model_name.is_some() {
1122            attributes.retain(|attribute| attribute.key.as_str() != "nemo_relay.model_name");
1123        }
1124        if self.otel_type == OpenTelemetryType::OpenInference && start_model_name.is_some() {
1125            super::openinference::remove_start_model_name(&mut attributes);
1126        }
1127        if self.otel_type != OpenTelemetryType::GenAi && is_trace_root {
1128            push_session_identity_attributes(&mut attributes, event);
1129        }
1130        let projected_attributes = if self.otel_type == OpenTelemetryType::GenAi {
1131            Vec::new()
1132        } else {
1133            attribute_mapping_inputs(&attributes, &self.attribute_mappings)
1134        };
1135        span.set_attributes(attributes);
1136        let span_context = local_parent_span_context(span.span_context());
1137        self.active_spans.insert(
1138            event.uuid(),
1139            ActiveSpan {
1140                span,
1141                span_context,
1142                start_model_name,
1143                projected_attributes,
1144                descendant_error_type: None,
1145                descendant_exception_type: None,
1146            },
1147        );
1148    }
1149
1150    fn process_end(&mut self, event: &Event) {
1151        let Some(mut active_span) = self.active_spans.remove(&event.uuid()) else {
1152            return;
1153        };
1154        self.record_completed_span_context(event.uuid(), active_span.span_context.clone());
1155
1156        super::set_span_status_from_event_metadata(&mut active_span.span, event);
1157        let mut attributes = match self.otel_type {
1158            OpenTelemetryType::Full => end_attributes(event),
1159            OpenTelemetryType::GenAi => super::otel_genai::end_attributes(event),
1160            OpenTelemetryType::OpenInference => super::openinference::end_attributes(event),
1161        };
1162        let is_error = metadata_string(event, "otel.status_code") == Some("ERROR");
1163        let explicit_error_type = metadata_string(event, "error.type");
1164        let error_type = is_error.then(|| {
1165            explicit_error_type
1166                .map(ToOwned::to_owned)
1167                .or(active_span.descendant_error_type.take())
1168                .unwrap_or_else(|| "_OTHER".to_string())
1169        });
1170        let exception_type = is_error
1171            .then(|| {
1172                metadata_string(event, "exception.type")
1173                    .map(ToOwned::to_owned)
1174                    .or(active_span.descendant_exception_type.take())
1175            })
1176            .flatten();
1177        if matches!(
1178            self.otel_type,
1179            OpenTelemetryType::Full | OpenTelemetryType::GenAi
1180        ) && let Some(error_type) = error_type.as_ref()
1181        {
1182            attributes.retain(|attribute| attribute.key.as_str() != "error.type");
1183            attributes.push(KeyValue::new("error.type", error_type.clone()));
1184        }
1185        if let Some(exception_type) = exception_type.as_ref() {
1186            active_span.span.add_event_with_timestamp(
1187                "exception",
1188                to_system_time(*event.timestamp()),
1189                vec![KeyValue::new("exception.type", exception_type.clone())],
1190            );
1191        }
1192        let end_model_name =
1193            model_name_for_llm_event(event).or_else(|| active_span.start_model_name.take());
1194        if self.otel_type == OpenTelemetryType::Full
1195            && let Some(model_name) = end_model_name.clone()
1196        {
1197            attributes.push(KeyValue::new("nemo_relay.model_name", model_name));
1198        }
1199        if self.otel_type == OpenTelemetryType::OpenInference
1200            && let Some(model_name) = end_model_name
1201        {
1202            super::openinference::push_model_name(&mut attributes, model_name);
1203        }
1204        if self.otel_type != OpenTelemetryType::GenAi && !self.attribute_mappings.is_empty() {
1205            let mut projected_attributes = active_span.projected_attributes;
1206            projected_attributes.extend(attributes.iter().cloned());
1207            attributes.extend(attribute_mapping_aliases(
1208                &projected_attributes,
1209                &self.attribute_mappings,
1210            ));
1211        }
1212        if is_error && let Some(parent_span) = self.find_parent_span_mut(event) {
1213            if parent_span.descendant_error_type.is_none() {
1214                parent_span.descendant_error_type = error_type;
1215            }
1216            if parent_span.descendant_exception_type.is_none() {
1217                parent_span.descendant_exception_type = exception_type;
1218            }
1219        }
1220        active_span.span.set_attributes(attributes);
1221        active_span
1222            .span
1223            .end_with_timestamp(to_system_time(*event.timestamp()));
1224    }
1225
1226    fn process_mark(&mut self, event: &Event) {
1227        if self.otel_type == OpenTelemetryType::GenAi {
1228            return;
1229        }
1230        if effective_mark_projection(event, self.mark_projection, &self.mark_exclude_names)
1231            == MarkProjection::Tool
1232        {
1233            self.process_mark_as_tool(event);
1234            return;
1235        }
1236        let mark_name = event.name().to_string();
1237        let timestamp = to_system_time(*event.timestamp());
1238        let mut attributes = self.mark_attributes(event);
1239        if event.name() == "session.start" {
1240            push_session_identity_attributes(&mut attributes, event);
1241        }
1242
1243        if self.find_parent_span(event).is_some() {
1244            apply_attribute_mappings(&mut attributes, &self.attribute_mappings);
1245            let parent_span = self
1246                .find_parent_span_mut(event)
1247                .expect("parent span was present during mark projection");
1248            parent_span
1249                .span
1250                .add_event_with_timestamp(mark_name, timestamp, attributes);
1251            return;
1252        }
1253
1254        let mut span = with_relay_ids(event.uuid(), || {
1255            self.tracer
1256                .span_builder(format!("mark:{mark_name}"))
1257                .with_kind(SpanKind::Internal)
1258                .with_start_time(timestamp)
1259                .start_with_context(&self.tracer, &self.parent_context(event))
1260        });
1261        if self.otel_type == OpenTelemetryType::OpenInference {
1262            super::openinference::push_orphan_mark_attributes(&mut attributes);
1263        } else {
1264            attributes.push(KeyValue::new("nemo_relay.mark.orphan", true));
1265        }
1266        apply_attribute_mappings(&mut attributes, &self.attribute_mappings);
1267        span.set_attributes(attributes);
1268        span.end_with_timestamp(timestamp);
1269    }
1270
1271    fn process_mark_as_tool(&mut self, event: &Event) {
1272        let timestamp = to_system_time(*event.timestamp());
1273        let orphan = self.find_parent_span(event).is_none();
1274        let mut attributes = self.mark_attributes(event);
1275        if event.name() == "session.start" {
1276            push_session_identity_attributes(&mut attributes, event);
1277        }
1278        attributes.push(KeyValue::new("nemo_relay.mark.projection", "tool"));
1279        if self.otel_type == OpenTelemetryType::OpenInference {
1280            super::openinference::push_tool_mark_attributes(&mut attributes, event);
1281        } else {
1282            attributes.push(KeyValue::new("nemo_relay.scope_type", "tool"));
1283        }
1284        if orphan {
1285            attributes.push(KeyValue::new("nemo_relay.mark.orphan", true));
1286        }
1287        apply_attribute_mappings(&mut attributes, &self.attribute_mappings);
1288
1289        let mut span = with_relay_ids(event.uuid(), || {
1290            self.tracer
1291                .span_builder(format!("mark:{}", event.name()))
1292                .with_kind(SpanKind::Internal)
1293                .with_start_time(timestamp)
1294                .start_with_context(&self.tracer, &self.parent_context(event))
1295        });
1296        span.set_attributes(attributes);
1297        span.end_with_timestamp(timestamp);
1298    }
1299
1300    fn mark_attributes(&self, event: &Event) -> Vec<KeyValue> {
1301        match self.otel_type {
1302            OpenTelemetryType::Full => mark_attributes(event),
1303            OpenTelemetryType::OpenInference => super::openinference::mark_attributes(event),
1304            OpenTelemetryType::GenAi => Vec::new(),
1305        }
1306    }
1307
1308    fn parent_context(&self, event: &Event) -> Context {
1309        if let Some(active_span) = self.find_parent_span(event) {
1310            return Context::new().with_remote_span_context(active_span.span_context.clone());
1311        }
1312        if let Some(span_context) = event
1313            .parent_uuid()
1314            .and_then(|uuid| self.completed_span_contexts.get(&uuid))
1315        {
1316            return Context::new().with_remote_span_context(span_context.clone());
1317        }
1318        let Some(parent_uuid) = event.parent_uuid() else {
1319            return Context::new();
1320        };
1321        let stack = current_scope_stack();
1322        let stack = stack.read().expect("scope stack lock poisoned");
1323        if !stack.is_propagated_parent(parent_uuid) {
1324            return Context::new();
1325        }
1326        let root_uuid = stack.root_uuid();
1327        Context::new().with_remote_span_context(SpanContext::new(
1328            relay_trace_id(root_uuid),
1329            relay_span_id(parent_uuid),
1330            TraceFlags::SAMPLED,
1331            true,
1332            TraceState::default(),
1333        ))
1334    }
1335
1336    fn parent_span_uuid(&self, event: &Event) -> Option<Uuid> {
1337        let parent_uuid = event.parent_uuid()?;
1338        self.active_spans
1339            .contains_key(&parent_uuid)
1340            .then_some(parent_uuid)
1341    }
1342
1343    fn find_parent_span(&self, event: &Event) -> Option<&ActiveSpan> {
1344        self.parent_span_uuid(event)
1345            .and_then(|uuid| self.active_spans.get(&uuid))
1346    }
1347
1348    fn find_parent_span_mut(&mut self, event: &Event) -> Option<&mut ActiveSpan> {
1349        self.parent_span_uuid(event)
1350            .and_then(|uuid| self.active_spans.get_mut(&uuid))
1351    }
1352
1353    fn remove_completed_span_context(&mut self, uuid: Uuid) {
1354        self.completed_span_contexts.remove(&uuid);
1355        self.completed_span_order
1356            .retain(|completed_uuid| *completed_uuid != uuid);
1357    }
1358
1359    fn record_completed_span_context(&mut self, uuid: Uuid, span_context: SpanContext) {
1360        if self
1361            .completed_span_contexts
1362            .insert(uuid, span_context)
1363            .is_none()
1364        {
1365            self.completed_span_order.push_back(uuid);
1366        }
1367        while self.completed_span_order.len() > COMPLETED_SPAN_CONTEXT_LIMIT {
1368            if let Some(expired) = self.completed_span_order.pop_front() {
1369                self.completed_span_contexts.remove(&expired);
1370            }
1371        }
1372    }
1373}
1374
1375fn metadata_string<'a>(event: &'a Event, key: &str) -> Option<&'a str> {
1376    event.metadata()?.get(key)?.as_str()
1377}
1378
1379fn span_kind(event: &Event) -> SpanKind {
1380    match semantic_scope_type(event) {
1381        Some(ScopeType::Llm) => SpanKind::Client,
1382        Some(
1383            ScopeType::Tool | ScopeType::Retriever | ScopeType::Embedder | ScopeType::Reranker,
1384        ) => SpanKind::Client,
1385        _ => SpanKind::Internal,
1386    }
1387}
1388
1389fn span_name(event: &Event) -> String {
1390    event.name().to_string()
1391}
1392
1393fn semantic_scope_type(event: &Event) -> Option<ScopeType> {
1394    event.scope_type()
1395}
1396
1397fn scope_type_name(scope_type: Option<ScopeType>) -> &'static str {
1398    match scope_type {
1399        Some(ScopeType::Agent) => "agent",
1400        Some(ScopeType::Function) => "function",
1401        Some(ScopeType::Tool) => "tool",
1402        Some(ScopeType::Llm) => "llm",
1403        Some(ScopeType::Retriever) => "retriever",
1404        Some(ScopeType::Embedder) => "embedder",
1405        Some(ScopeType::Reranker) => "reranker",
1406        Some(ScopeType::Guardrail) => "guardrail",
1407        Some(ScopeType::Evaluator) => "evaluator",
1408        Some(ScopeType::Custom) => "custom",
1409        Some(ScopeType::Unknown) | None => "unknown",
1410    }
1411}
1412
1413fn start_attributes(event: &Event) -> Vec<KeyValue> {
1414    let mut attributes = common_attributes(event);
1415    push_serialized_top_level_attributes(
1416        &mut attributes,
1417        "nemo_relay.handle_attributes",
1418        event.attributes(),
1419    );
1420    push_top_level_json_attributes(&mut attributes, "nemo_relay.start.data", event.data());
1421    push_top_level_json_attributes(
1422        &mut attributes,
1423        "nemo_relay.start.metadata",
1424        event.metadata(),
1425    );
1426    push_top_level_json_attributes(&mut attributes, "nemo_relay.start.input", event.input());
1427    attributes
1428}
1429
1430fn end_attributes(event: &Event) -> Vec<KeyValue> {
1431    let mut attributes = Vec::new();
1432    push_top_level_json_attributes(&mut attributes, "nemo_relay.end.data", event.data());
1433    push_top_level_json_attributes(&mut attributes, "nemo_relay.end.metadata", event.metadata());
1434    push_top_level_json_attributes(&mut attributes, "nemo_relay.end.output", event.output());
1435    if event
1436        .category()
1437        .is_some_and(|category| category.as_str() == "llm")
1438        && let Some((cost, currency)) = cost_from_llm_event(event)
1439    {
1440        attributes.push(KeyValue::new("nemo_relay.llm.cost.total", cost));
1441        attributes.push(KeyValue::new("nemo_relay.llm.cost.currency", currency));
1442    }
1443    if let Some(response) = event.annotated_response()
1444        && let Some(summary) = response.optimization_summary.as_ref()
1445    {
1446        push_optimization_attributes(&mut attributes, summary);
1447    }
1448    attributes
1449}
1450
1451fn push_optimization_attributes(
1452    attributes: &mut Vec<KeyValue>,
1453    summary: &crate::codec::optimization::LlmOptimizationSummary,
1454) {
1455    crate::observability::push_common_optimization_attributes(attributes, summary);
1456}
1457
1458fn cost_from_llm_event(event: &Event) -> Option<(f64, String)> {
1459    if let Some(response) = event.normalized_llm_response() {
1460        let response = response.as_ref();
1461        if let Some(usage) = response.usage.as_ref() {
1462            if let Some(cost) = usage.cost.as_ref() {
1463                return cost_total_and_currency(cost);
1464            }
1465            if let Some(cost) = estimate_cost_for_response_or_requested_model(
1466                event,
1467                response.model.as_deref(),
1468                usage,
1469            ) {
1470                return cost_total_and_currency(&cost);
1471            }
1472        }
1473    }
1474    if let Some(cost) =
1475        manual::cost_from_manual_llm_output(event.output(), manual::ManualCostPolicy::AnyCurrency)
1476    {
1477        return Some(cost);
1478    }
1479    let usage = manual::usage_from_manual_llm_output(event.output())?;
1480    estimate_cost_for_response_or_model(
1481        Some(event.name()),
1482        event.model_name(),
1483        manual::model_name_from_manual_llm_output(event.output()),
1484        &usage,
1485    )
1486    .and_then(|cost| cost_total_and_currency(&cost))
1487}
1488
1489fn cost_total_and_currency(cost: &CostEstimate) -> Option<(f64, String)> {
1490    Some((cost.total_or_component_sum()?, cost.currency.clone()))
1491}
1492
1493fn mark_attributes(event: &Event) -> Vec<KeyValue> {
1494    let mut attributes = vec![
1495        KeyValue::new("nemo_relay.mark.uuid", event.uuid().to_string()),
1496        KeyValue::new(
1497            "nemo_relay.mark.parent_uuid",
1498            event
1499                .parent_uuid()
1500                .map(|uuid| uuid.to_string())
1501                .unwrap_or_default(),
1502        ),
1503    ];
1504    push_serialized_top_level_attributes(
1505        &mut attributes,
1506        "nemo_relay.mark.attributes",
1507        event.attributes(),
1508    );
1509    push_top_level_json_attributes(&mut attributes, "nemo_relay.mark.data", event.data());
1510    push_top_level_json_attributes(
1511        &mut attributes,
1512        "nemo_relay.mark.metadata",
1513        event.metadata(),
1514    );
1515    if let Some(category) = event.category() {
1516        attributes.push(KeyValue::new(
1517            "nemo_relay.mark.category",
1518            category.as_str().to_string(),
1519        ));
1520    }
1521    push_serialized_top_level_attributes(
1522        &mut attributes,
1523        "nemo_relay.mark.category_profile",
1524        event.category_profile(),
1525    );
1526    attributes
1527}
1528
1529fn common_attributes(event: &Event) -> Vec<KeyValue> {
1530    let mut attributes = vec![
1531        KeyValue::new("nemo_relay.uuid", event.uuid().to_string()),
1532        KeyValue::new(
1533            "nemo_relay.parent_uuid",
1534            event
1535                .parent_uuid()
1536                .map(|uuid| uuid.to_string())
1537                .unwrap_or_default(),
1538        ),
1539        KeyValue::new(
1540            "nemo_relay.scope_type",
1541            scope_type_name(semantic_scope_type(event)),
1542        ),
1543    ];
1544
1545    if let Some(model_name) = model_name_for_llm_event(event) {
1546        attributes.push(KeyValue::new("nemo_relay.model_name", model_name));
1547    }
1548    if let Some(tool_call_id) = event.tool_call_id() {
1549        attributes.push(KeyValue::new(
1550            "nemo_relay.tool_call_id",
1551            tool_call_id.to_string(),
1552        ));
1553    }
1554
1555    attributes
1556}
1557
1558fn local_parent_span_context(span_context: &SpanContext) -> SpanContext {
1559    SpanContext::new(
1560        span_context.trace_id(),
1561        span_context.span_id(),
1562        span_context.trace_flags(),
1563        false,
1564        span_context.trace_state().clone(),
1565    )
1566}
1567
1568fn to_system_time(timestamp: DateTime<Utc>) -> SystemTime {
1569    let seconds = timestamp.timestamp();
1570    let nanos = timestamp.timestamp_subsec_nanos();
1571    if seconds >= 0 {
1572        UNIX_EPOCH + Duration::new(seconds as u64, nanos)
1573    } else if nanos == 0 {
1574        UNIX_EPOCH - Duration::new(seconds.unsigned_abs(), 0)
1575    } else {
1576        UNIX_EPOCH - Duration::new(seconds.unsigned_abs() - 1, 1_000_000_000 - nanos)
1577    }
1578}
1579
1580#[cfg(test)]
1581#[path = "../../tests/unit/observability/otel_tests.rs"]
1582mod tests;