diff --git a/crates/core/src/observability/mod.rs b/crates/core/src/observability/mod.rs index 8d14c24df..128379fad 100644 --- a/crates/core/src/observability/mod.rs +++ b/crates/core/src/observability/mod.rs @@ -72,6 +72,18 @@ pub(crate) fn default_mark_exclude_names() -> Vec { vec!["llm.chunk".to_string()] } +#[cfg(any(feature = "otel", feature = "openinference"))] +pub(crate) fn relay_trace_id(uuid: uuid::Uuid) -> opentelemetry::trace::TraceId { + opentelemetry::trace::TraceId::from_bytes(*uuid.as_bytes()) +} + +#[cfg(any(feature = "otel", feature = "openinference"))] +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) +} + #[cfg(any(feature = "otel", feature = "openinference"))] pub(crate) fn push_common_optimization_attributes( attributes: &mut Vec, @@ -561,3 +573,17 @@ where #[cfg(all(test, any(feature = "otel", feature = "openinference")))] #[path = "../../tests/unit/observability/attribute_projection_tests.rs"] mod attribute_projection_tests; + +#[cfg(all(test, any(feature = "otel", feature = "openinference")))] +mod tests { + use super::{relay_span_id, relay_trace_id}; + use uuid::Uuid; + + #[test] + fn relay_id_conversions_preserve_zero_bytes() { + let uuid = Uuid::nil(); + + assert_eq!(relay_trace_id(uuid).to_bytes(), [0; 16]); + assert_eq!(relay_span_id(uuid).to_bytes(), [0; 8]); + } +} diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 709d3d114..b8cc19ae6 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -25,7 +25,8 @@ use super::{ attribute_mapping_inputs, default_mark_exclude_names, effective_mark_projection, estimate_cost_for_response_or_model, estimate_cost_for_response_or_requested_model, manual, merge_usage, model_name_for_llm_event, push_serialized_top_level_attributes, - push_session_identity_attributes, push_top_level_json_attributes, validate_attribute_mappings, + push_session_identity_attributes, push_top_level_json_attributes, relay_span_id, + relay_trace_id, validate_attribute_mappings, }; use crate::api::event::{Event, EventNormalizationExt, ScopeCategory}; use crate::api::runtime::EventSubscriberFn; @@ -634,6 +635,8 @@ impl OpenInferenceEventProcessor { .span_builder(span_name(event)) .with_kind(span_kind(event)) .with_start_time(to_system_time(*event.timestamp())) + .with_trace_id(relay_trace_id(event.uuid())) + .with_span_id(relay_span_id(event.uuid())) .start_with_context(&self.tracer, &parent_context); let mut attributes = start_attributes(event); if is_trace_root { @@ -703,6 +706,8 @@ impl OpenInferenceEventProcessor { .span_builder(format!("mark:{mark_name}")) .with_kind(SpanKind::Internal) .with_start_time(timestamp) + .with_trace_id(relay_trace_id(event.uuid())) + .with_span_id(relay_span_id(event.uuid())) .start_with_context(&self.tracer, &self.parent_context(event)); attributes.push(KeyValue::new( oi::OPENINFERENCE_SPAN_KIND, @@ -737,6 +742,8 @@ impl OpenInferenceEventProcessor { .span_builder(format!("mark:{}", event.name())) .with_kind(SpanKind::Internal) .with_start_time(timestamp) + .with_trace_id(relay_trace_id(event.uuid())) + .with_span_id(relay_span_id(event.uuid())) .start_with_context(&self.tracer, &self.parent_context(event)); span.set_attributes(attributes); span.end_with_timestamp(timestamp); diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index 97a493db9..9b3e8c929 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -25,7 +25,8 @@ use super::{ attribute_mapping_inputs, default_mark_exclude_names, effective_mark_projection, estimate_cost_for_response_or_model, estimate_cost_for_response_or_requested_model, manual, model_name_for_llm_event, push_serialized_top_level_attributes, - push_session_identity_attributes, push_top_level_json_attributes, validate_attribute_mappings, + push_session_identity_attributes, push_top_level_json_attributes, relay_span_id, + relay_trace_id, validate_attribute_mappings, }; use crate::api::event::{Event, EventNormalizationExt, ScopeCategory}; use crate::api::runtime::EventSubscriberFn; @@ -627,6 +628,8 @@ impl OtelEventProcessor { .span_builder(span_name(event)) .with_kind(span_kind(event)) .with_start_time(to_system_time(*event.timestamp())) + .with_trace_id(relay_trace_id(event.uuid())) + .with_span_id(relay_span_id(event.uuid())) .start_with_context(&self.tracer, &parent_context); let mut attributes = start_attributes(event); if is_trace_root { @@ -697,6 +700,8 @@ impl OtelEventProcessor { .span_builder(format!("mark:{mark_name}")) .with_kind(SpanKind::Internal) .with_start_time(timestamp) + .with_trace_id(relay_trace_id(event.uuid())) + .with_span_id(relay_span_id(event.uuid())) .start_with_context(&self.tracer, &self.parent_context(event)); attributes.push(KeyValue::new("nemo_relay.mark.orphan", true)); apply_attribute_mappings(&mut attributes, &self.attribute_mappings); @@ -723,6 +728,8 @@ impl OtelEventProcessor { .span_builder(format!("mark:{}", event.name())) .with_kind(SpanKind::Internal) .with_start_time(timestamp) + .with_trace_id(relay_trace_id(event.uuid())) + .with_span_id(relay_span_id(event.uuid())) .start_with_context(&self.tracer, &self.parent_context(event)); span.set_attributes(attributes); span.end_with_timestamp(timestamp); diff --git a/crates/core/tests/unit/observability/exporter_parity_tests.rs b/crates/core/tests/unit/observability/exporter_parity_tests.rs index ea40a4f19..832fe5fcc 100644 --- a/crates/core/tests/unit/observability/exporter_parity_tests.rs +++ b/crates/core/tests/unit/observability/exporter_parity_tests.rs @@ -819,7 +819,7 @@ fn test_manual_fallback_payload_parity() { } #[test] -fn session_instance_correlates_distinct_otel_and_openinference_traces() { +fn session_instance_correlates_otel_and_openinference_traces() { let scope_stack = create_scope_stack(); let expected_root_uuid = scope_stack.read().unwrap().root_uuid().to_string(); let uuid = Uuid::now_v7(); @@ -874,7 +874,7 @@ fn session_instance_correlates_distinct_otel_and_openinference_traces() { openinference_attributes["nemo_relay.session.instance_id"], expected_root_uuid ); - assert_ne!( + assert_eq!( otel.span_context.trace_id(), openinference.span_context.trace_id() ); diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index 6f3cfdc0d..25fca6174 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -2361,13 +2361,17 @@ fn tool_semantic_names_exist_without_input_payload() { } #[test] -fn preserves_parent_child_relationships() { +fn derives_span_ids_from_relay_uuids() { let (provider, exporter) = make_provider(); - let mut processor = - OpenInferenceEventProcessor::new(provider.clone(), "test-scope".to_string()); + let mut processor = OpenInferenceEventProcessor::new_with_mark_projection( + provider.clone(), + "test-scope".to_string(), + MarkProjection::Tool, + ); - let root_uuid = Uuid::now_v7(); - let child_uuid = Uuid::now_v7(); + let root_uuid = Uuid::from_u128(0x018f_0f0f_0f0f_7000_8123_4567_89ab_cdef); + let child_uuid = Uuid::from_u128(0x018f_0f0f_0f10_7000_8fed_cba9_8765_4321); + let orphan_uuid = Uuid::from_u128(0x018f_0f0f_0f11_7000_8011_2233_4455_6677); processor.process(&make_start_event( root_uuid, @@ -2397,11 +2401,19 @@ fn preserves_parent_child_relationships() { ScopeType::Agent, None, )); + processor.process(&Event::Mark(MarkEvent::new( + BaseEvent::builder() + .uuid(orphan_uuid) + .name("checkpoint") + .build(), + None, + None, + ))); processor.force_flush().unwrap(); let spans = exporter.get_finished_spans().unwrap(); - assert_eq!(spans.len(), 2); + assert_eq!(spans.len(), 3); let parent = spans .iter() .find(|span| span.name.as_ref() == "agent") @@ -2410,13 +2422,37 @@ fn preserves_parent_child_relationships() { .iter() .find(|span| span.name.as_ref() == "model-call") .unwrap(); + let orphan = spans + .iter() + .find(|span| span.name.as_ref() == "mark:checkpoint") + .unwrap(); + assert_eq!( + parent.span_context.trace_id().to_bytes(), + *root_uuid.as_bytes() + ); + assert_eq!( + parent.span_context.span_id().to_bytes(), + root_uuid.as_bytes()[8..] + ); assert_eq!( child.span_context.trace_id(), parent.span_context.trace_id() ); + assert_eq!( + child.span_context.span_id().to_bytes(), + child_uuid.as_bytes()[8..] + ); assert_eq!(child.parent_span_id, parent.span_context.span_id()); assert!(!child.parent_span_is_remote); + assert_eq!( + orphan.span_context.trace_id().to_bytes(), + *orphan_uuid.as_bytes() + ); + assert_eq!( + orphan.span_context.span_id().to_bytes(), + orphan_uuid.as_bytes()[8..] + ); } #[test] diff --git a/crates/core/tests/unit/observability/otel_tests.rs b/crates/core/tests/unit/observability/otel_tests.rs index 2422e43a9..bbdcd082c 100644 --- a/crates/core/tests/unit/observability/otel_tests.rs +++ b/crates/core/tests/unit/observability/otel_tests.rs @@ -1049,12 +1049,17 @@ fn records_span_start_mark_and_end() { } #[test] -fn preserves_parent_child_relationships() { +fn derives_span_ids_from_relay_uuids() { let (provider, exporter) = make_provider(); - let mut processor = OtelEventProcessor::new(provider.clone(), "test-scope".to_string()); + let mut processor = OtelEventProcessor::new_with_mark_projection( + provider.clone(), + "test-scope".to_string(), + MarkProjection::Tool, + ); - let root_uuid = Uuid::now_v7(); - let child_uuid = Uuid::now_v7(); + let root_uuid = Uuid::from_u128(0x018f_0f0f_0f0f_7000_8123_4567_89ab_cdef); + let child_uuid = Uuid::from_u128(0x018f_0f0f_0f10_7000_8fed_cba9_8765_4321); + let orphan_uuid = Uuid::from_u128(0x018f_0f0f_0f11_7000_8011_2233_4455_6677); processor.process(&make_start_event( root_uuid, @@ -1084,11 +1089,19 @@ fn preserves_parent_child_relationships() { ScopeType::Agent, None, )); + processor.process(&Event::Mark(MarkEvent::new( + BaseEvent::builder() + .uuid(orphan_uuid) + .name("checkpoint") + .build(), + None, + None, + ))); processor.force_flush().unwrap(); let spans = exporter.get_finished_spans().unwrap(); - assert_eq!(spans.len(), 2); + assert_eq!(spans.len(), 3); let parent = spans .iter() .find(|span| span.name.as_ref() == "agent") @@ -1097,13 +1110,37 @@ fn preserves_parent_child_relationships() { .iter() .find(|span| span.name.as_ref() == "model-call") .unwrap(); + let orphan = spans + .iter() + .find(|span| span.name.as_ref() == "mark:checkpoint") + .unwrap(); + assert_eq!( + parent.span_context.trace_id().to_bytes(), + *root_uuid.as_bytes() + ); + assert_eq!( + parent.span_context.span_id().to_bytes(), + root_uuid.as_bytes()[8..] + ); assert_eq!( child.span_context.trace_id(), parent.span_context.trace_id() ); + assert_eq!( + child.span_context.span_id().to_bytes(), + child_uuid.as_bytes()[8..] + ); assert_eq!(child.parent_span_id, parent.span_context.span_id()); assert!(!child.parent_span_is_remote); + assert_eq!( + orphan.span_context.trace_id().to_bytes(), + *orphan_uuid.as_bytes() + ); + assert_eq!( + orphan.span_context.span_id().to_bytes(), + orphan_uuid.as_bytes()[8..] + ); } #[test]