codex: address PR review feedback (#18725)

codex-rs/app-server/src/otel_reload.rs

@@ -116,6 +116,7 @@ impl OtelReloader {
     pub(crate) fn shutdown(&self) {
         self.logger_layer.replace_provider(/*provider*/ None);
         self.trace_layer.shutdown();
+        OtelProvider::replace_global_metrics(/*provider*/ None);
         let mut state = self
             .state
             .lock()
@@ -193,6 +194,7 @@ impl Drop for OtelReloadCommit {
             state.key = self.previous_key.clone();
             self.logger_layer.replace_provider(state.provider.as_ref());
             self.trace_layer.replace_provider(state.provider.as_ref());
+            OtelProvider::replace_global_metrics(state.provider.as_ref());
         } else if let Some(old_provider) = self.old_provider.take() {
             state.retired_providers.push(old_provider);
         }

codex-rs/otel/src/logger_layer.rs

@@ -10,7 +10,6 @@ use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge;
 use opentelemetry_sdk::logs::SdkLoggerProvider;
 use std::sync::Arc;
 use std::sync::RwLock;
-use tracing::subscriber::Interest;
 use tracing_subscriber::Layer;
 use tracing_subscriber::layer::Context;
 use tracing_subscriber::registry::LookupSpan;
@@ -45,25 +44,6 @@ impl<S> Layer<S> for OtelLoggerLayer
 where
     S: tracing::Subscriber + for<'span> LookupSpan<'span>,
 {
-    fn register_callsite(&self, metadata: &'static tracing::Metadata<'static>) -> Interest {
-        if OtelProvider::log_export_filter(metadata) {
-            Interest::sometimes()
-        } else {
-            Interest::never()
-        }
-    }
-
-    fn enabled(&self, metadata: &tracing::Metadata<'_>, _ctx: Context<'_, S>) -> bool {
-        if !OtelProvider::log_export_filter(metadata) {
-            return false;
-        }
-
-        self.bridge
-            .read()
-            .unwrap_or_else(std::sync::PoisonError::into_inner)
-            .is_some()
-    }
-
     fn on_event(&self, event: &tracing::Event<'_>, ctx: Context<'_, S>) {
         if !OtelProvider::log_export_filter(event.metadata()) {
             return;
@@ -84,3 +64,48 @@ fn logger_bridge(provider: Option<&OtelProvider>) -> Option<LoggerBridge> {
         .and_then(|provider| provider.logger.as_ref())
         .map(OpenTelemetryTracingBridge::new)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::Mutex;
+    use tracing::Subscriber;
+    use tracing_subscriber::prelude::*;
+
+    #[derive(Clone, Default)]
+    struct CountingLayer {
+        events: Arc<Mutex<usize>>,
+    }
+
+    impl<S> Layer<S> for CountingLayer
+    where
+        S: Subscriber,
+    {
+        fn on_event(&self, _event: &tracing::Event<'_>, _ctx: Context<'_, S>) {
+            let mut events = self
+                .events
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            *events += 1;
+        }
+    }
+
+    #[test]
+    fn non_otel_events_continue_to_reach_other_layers() {
+        let counting_layer = CountingLayer::default();
+        let events = Arc::clone(&counting_layer.events);
+        let subscriber = tracing_subscriber::registry()
+            .with(counting_layer)
+            .with(OtelLoggerLayer::default());
+        let _guard = tracing::subscriber::set_default(subscriber);
+
+        tracing::info!(target: "codex_core::unrelated", "visible to other layers");
+
+        assert_eq!(
+            *events
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner),
+            1
+        );
+    }
+}

codex-rs/otel/src/provider.rs

@@ -156,6 +156,10 @@ impl OtelProvider {
     pub fn metrics(&self) -> Option<&MetricsClient> {
         self.metrics.as_ref()
     }
+
+    pub fn replace_global_metrics(provider: Option<&Self>) {
+        crate::metrics::replace_global(provider.and_then(|provider| provider.metrics.clone()));
+    }
 }
 
 impl Drop for OtelProvider {

codex-rs/otel/src/trace_layer.rs

@@ -162,7 +162,15 @@ where
             return;
         }
 
-        if let Some(layer) = self.current_layer() {
+        let span_layer = event
+            .parent()
+            .and_then(|id| self.layer_for_span(id))
+            .or_else(|| {
+                ctx.event_span(event)
+                    .and_then(|span| self.layer_for_span(&span.id()))
+            });
+
+        if let Some(layer) = span_layer.or_else(|| self.current_layer()) {
             layer.on_event(event, ctx);
         }
     }
@@ -328,4 +336,26 @@ mod tests {
             0
         );
     }
+
+    #[test]
+    fn span_events_use_original_layer_after_provider_replacement() {
+        let (first_provider, first_spans) = provider_with_tracer("first");
+        let (layer, handle) = OtelTraceLayer::from_provider(Some(&first_provider));
+        let subscriber = tracing_subscriber::registry().with(layer);
+        let _guard = tracing::subscriber::set_default(subscriber);
+
+        let span = trace_span!("old_span");
+        {
+            let _entered = span.enter();
+            handle.replace_provider(/*provider*/ None);
+            tracing::info!(target: "codex_otel.trace_safe.test", "old_span_event");
+        }
+        drop(span);
+
+        let spans = first_spans
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        assert_eq!(spans.len(), 1);
+        assert_eq!(spans[0].events.len(), 1);
+    }
 }