fix(audit): address OTLP payload review feedback

Signed-off-by: Yuting Wu (DLAlgo) <yutwu@nvidia.com>

Author: nvyutwu

lib/llm/src/audit/config.rs

@@ -4,6 +4,9 @@
 use std::sync::OnceLock;
 use std::sync::atomic::{AtomicU8, Ordering};
 
+#[cfg(test)]
+use std::sync::Mutex;
+
 use dynamo_runtime::config::environment_names::llm::audit as env_audit;
 
 use crate::telemetry::parse_sink_names;
@@ -35,6 +38,9 @@ pub struct AuditPolicy {
 static POLICY: OnceLock<AuditPolicy> = OnceLock::new();
 static CAPTURE_STATE: AtomicU8 = AtomicU8::new(CAPTURE_UNINITIALIZED);
 
+#[cfg(test)]
+static TEST_POLICY_OVERRIDE: Mutex<Option<&'static AuditPolicy>> = Mutex::new(None);
+
 /// Audit is enabled if we have at least one sink
 fn load_from_env() -> AuditPolicy {
     let sinks = std::env::var(env_audit::DYN_AUDIT_SINKS)
@@ -91,9 +97,33 @@ fn load_from_env() -> AuditPolicy {
 }
 
 pub fn policy() -> &'static AuditPolicy {
+    #[cfg(test)]
+    if let Some(policy) = *TEST_POLICY_OVERRIDE
+        .lock()
+        .expect("test policy lock poisoned")
+    {
+        return policy;
+    }
+
     POLICY.get_or_init(load_from_env)
 }
 
+#[cfg(test)]
+pub(crate) fn override_policy_from_env_for_test() {
+    let policy = Box::leak(Box::new(load_from_env()));
+    *TEST_POLICY_OVERRIDE
+        .lock()
+        .expect("test policy lock poisoned") = Some(policy);
+}
+
+#[cfg(test)]
+pub(crate) fn clear_policy_override_for_test() {
+    *TEST_POLICY_OVERRIDE
+        .lock()
+        .expect("test policy lock poisoned") = None;
+    mark_capture_inactive();
+}
+
 pub(crate) fn mark_capture_active() {
     CAPTURE_STATE.store(CAPTURE_ACTIVE, Ordering::Release);
 }
@@ -102,7 +132,7 @@ pub(crate) fn mark_capture_inactive() {
     CAPTURE_STATE.store(CAPTURE_INACTIVE, Ordering::Release);
 }
 
-pub fn capture_enabled() -> bool {
+pub(crate) fn capture_enabled() -> bool {
     // Require the explicit ACTIVE transition so that publishes happening before
     // `init_from_env_with_shutdown` finishes — i.e. while the bus is still
     // uninitialized — are skipped at the `create_handle` gate rather than
@@ -114,7 +144,7 @@ pub fn capture_enabled() -> bool {
     policy.enabled && CAPTURE_STATE.load(Ordering::Acquire) == CAPTURE_ACTIVE
 }
 
-pub fn otel_sink_capture_enabled() -> bool {
+pub(crate) fn otel_sink_capture_enabled() -> bool {
     let policy = policy();
     policy.enabled
         && CAPTURE_STATE.load(Ordering::Acquire) == CAPTURE_ACTIVE

lib/llm/src/audit/handle.rs

@@ -31,9 +31,9 @@ impl AuditHttpRequestHeaders {
 /// Distinguishes the two record types emitted per chat completion.
 ///
 /// Request and response are published as separate `AuditRecord`s sharing the same
-/// `request_id`. Downstream consumers correlate by `request_id`; the request record
-/// is emitted before the worker dispatches, the response record is emitted after the
-/// response stream completes successfully. On client cancel mid-stream (or
+/// `request_id`. Downstream consumers correlate by `request_id`; the request emit
+/// is scheduled before the worker dispatches, the response record is emitted after
+/// the response stream completes successfully. On client cancel mid-stream (or
 /// aggregation failure) only the request record is emitted.
 #[derive(Serialize, Deserialize, Clone, Copy, Debug, Eq, PartialEq)]
 #[serde(rename_all = "lowercase")]
@@ -74,9 +74,10 @@ impl AuditHandle {
         &self.request_id
     }
 
-    /// Publish a `Request` event record on the audit bus. Call once, as soon as the
-    /// request is captured and before worker dispatch — this lets downstream
-    /// observers see hung / canceled requests that never produce a response record.
+    /// Publish a `Request` event record on the audit bus. Call once after the
+    /// request is captured. The preprocessor schedules this before worker
+    /// dispatch so downstream observers can see hung / canceled requests that
+    /// never produce a response record.
     pub fn emit_request(&self, request: Arc<NvCreateChatCompletionRequest>) {
         let rec = AuditRecord {
             schema_version: 1,
@@ -183,20 +184,28 @@ mod tests {
         serde_json::from_value(json).expect("Failed to create test response")
     }
 
+    struct AuditPolicyResetGuard;
+
+    impl Drop for AuditPolicyResetGuard {
+        fn drop(&mut self) {
+            crate::audit::config::clear_policy_override_for_test();
+        }
+    }
+
     /// Test that DYN_AUDIT_FORCE_LOGGING=true bypasses store=false
     /// When force logging is enabled, audit handle should be created even when store=false
     #[test]
+    #[serial_test::serial]
     fn test_force_logging_bypasses_store() {
         with_vars(
             [
                 ("DYN_AUDIT_SINKS", Some("stderr")),
                 ("DYN_AUDIT_FORCE_LOGGING", Some("true")),
             ],
             || {
-                // `capture_enabled()` now requires `CAPTURE_ACTIVE`; mimic the
-                // audit init lifecycle (`init_from_env_with_shutdown`) instead of
-                // relying on the old "uninitialized counts as enabled" semantics.
+                crate::audit::config::override_policy_from_env_for_test();
                 crate::audit::config::mark_capture_active();
+                let _reset_guard = AuditPolicyResetGuard;
 
                 let request = create_test_request("test-model", false);
                 let handle = create_handle(&request, "test-id", None);

lib/llm/src/audit/otel_sink.rs

@@ -5,7 +5,9 @@
 //!
 //! Emits exactly one OTLP `LogRecord` per `AuditRecord`. The exporter is
 //! constructed once at sink init (not per emit). Network I/O happens on the
-//! SDK's internal batch processor; `emit()` is non-blocking enqueue.
+//! SDK's internal batch processor; `emit()` is non-blocking enqueue. The audit
+//! worker calls `force_flush()` after draining on shutdown, but abrupt process
+//! teardown can still lose buffered OTLP records.
 //!
 //! Transport follows `OTEL_EXPORTER_OTLP_LOGS_PROTOCOL` with
 //! `OTEL_EXPORTER_OTLP_PROTOCOL` as fallback. Supported values are
@@ -22,6 +24,7 @@ use axum::http::HeaderValue;
 use dynamo_runtime::config::environment_names::{
     llm::audit as env_audit, logging::otlp as env_otlp,
 };
+use opentelemetry::Context;
 use opentelemetry::logs::{AnyValue, LogRecord, Logger, LoggerProvider, Severity};
 use opentelemetry_otlp::{Protocol, WithExportConfig};
 use opentelemetry_sdk::Resource;
@@ -79,11 +82,8 @@ const AUDIT_INSTRUMENTATION_SCOPE: &str = "dynamo.payload";
 const DEFAULT_SERVICE_NAME: &str = "dynamo";
 
 pub struct OtelSink {
-    /// Held so the SDK's batch processor flushes when the sink is dropped on
-    /// audit-bus shutdown. The field is never read directly — its job is to
-    /// keep the provider alive for the sink's lifetime. TODO(phase D): wire
-    /// an explicit `force_flush` hook on the worker cancellation path so
-    /// records aren't lost if the runtime is torn down before Drop runs.
+    /// Held so the SDK's batch processor stays alive for the sink's lifetime
+    /// and can be force-flushed when the audit worker shuts down.
     #[allow(dead_code)]
     provider: SdkLoggerProvider,
     logger: SdkLogger,
@@ -175,6 +175,24 @@ impl OtlpLogsProtocol {
     }
 }
 
+fn logs_endpoint_from_env(protocol: OtlpLogsProtocol) -> String {
+    if let Ok(endpoint) = std::env::var(env_otlp::OTEL_EXPORTER_OTLP_LOGS_ENDPOINT) {
+        return endpoint;
+    }
+
+    if let Ok(endpoint) = std::env::var(env_otlp::OTEL_EXPORTER_OTLP_ENDPOINT) {
+        return match protocol {
+            OtlpLogsProtocol::HttpProtobuf => {
+                let trimmed = endpoint.trim_end_matches('/');
+                format!("{trimmed}/v1/logs")
+            }
+            OtlpLogsProtocol::Grpc => endpoint,
+        };
+    }
+
+    protocol.default_endpoint().to_string()
+}
+
 fn render_header_value(name: &str, value: &HeaderValue, policy: &OtelHeaderPolicy) -> String {
     if policy.should_redact(name) {
         return REDACTED_HEADER_VALUE.to_string();
@@ -267,9 +285,7 @@ impl OtelSink {
 
     pub async fn from_policy(policy: &AuditPolicy) -> Result<Self> {
         let protocol = OtlpLogsProtocol::from_env();
-        let endpoint = std::env::var(env_otlp::OTEL_EXPORTER_OTLP_LOGS_ENDPOINT)
-            .or_else(|_| std::env::var(env_otlp::OTEL_EXPORTER_OTLP_TRACES_ENDPOINT))
-            .unwrap_or_else(|_| protocol.default_endpoint().to_string());
+        let endpoint = logs_endpoint_from_env(protocol);
 
         let exporter = match protocol {
             OtlpLogsProtocol::HttpProtobuf => opentelemetry_otlp::LogExporter::builder()
@@ -468,8 +484,23 @@ impl AuditSink for OtelSink {
             record.add_attribute("audit_drop_reason", AnyValue::String(reason.into()));
         }
         record.add_attribute("payload", AnyValue::String(payload.into()));
+
+        // Audit OTLP export is an explicit sink, not telemetry generated while
+        // exporting telemetry. Use a fresh context so a globally suppressed
+        // tracing bridge cannot cause the direct LogRecord emit to be skipped.
+        let _guard = Context::new().attach();
         self.logger.emit(record);
     }
+
+    async fn shutdown(&self) {
+        if let Err(err) = self.provider.force_flush() {
+            tracing::warn!(
+                target: "dynamo_llm::audit",
+                error = %err,
+                "audit otel: force_flush failed during shutdown"
+            );
+        }
+    }
 }
 
 #[cfg(test)]
@@ -815,4 +846,59 @@ mod tests {
             || assert_eq!(OtlpLogsProtocol::from_env(), OtlpLogsProtocol::Grpc),
         );
     }
+
+    #[test]
+    #[serial]
+    fn logs_endpoint_uses_signal_specific_endpoint_first() {
+        temp_env::with_vars(
+            [
+                (
+                    env_otlp::OTEL_EXPORTER_OTLP_LOGS_ENDPOINT,
+                    Some("http://collector:9999/custom/logs"),
+                ),
+                (
+                    env_otlp::OTEL_EXPORTER_OTLP_ENDPOINT,
+                    Some("http://collector:4318"),
+                ),
+                (
+                    env_otlp::OTEL_EXPORTER_OTLP_TRACES_ENDPOINT,
+                    Some("http://collector:4317/v1/traces"),
+                ),
+            ],
+            || {
+                assert_eq!(
+                    logs_endpoint_from_env(OtlpLogsProtocol::HttpProtobuf),
+                    "http://collector:9999/custom/logs"
+                );
+            },
+        );
+    }
+
+    #[test]
+    #[serial]
+    fn logs_endpoint_falls_back_to_generic_endpoint_not_traces_endpoint() {
+        temp_env::with_vars(
+            [
+                (env_otlp::OTEL_EXPORTER_OTLP_LOGS_ENDPOINT, None::<&str>),
+                (
+                    env_otlp::OTEL_EXPORTER_OTLP_ENDPOINT,
+                    Some("http://collector:4318"),
+                ),
+                (
+                    env_otlp::OTEL_EXPORTER_OTLP_TRACES_ENDPOINT,
+                    Some("http://collector:4317/v1/traces"),
+                ),
+            ],
+            || {
+                assert_eq!(
+                    logs_endpoint_from_env(OtlpLogsProtocol::HttpProtobuf),
+                    "http://collector:4318/v1/logs"
+                );
+                assert_eq!(
+                    logs_endpoint_from_env(OtlpLogsProtocol::Grpc),
+                    "http://collector:4318"
+                );
+            },
+        );
+    }
 }

lib/llm/src/audit/sink.rs

@@ -1,6 +1,7 @@
 // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+use std::io::Write as _;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -26,6 +27,7 @@ use super::{
 pub trait AuditSink: Send + Sync {
     fn name(&self) -> &'static str;
     async fn emit(&self, rec: &AuditRecord);
+    async fn shutdown(&self) {}
 }
 
 pub struct StderrSink;
@@ -37,7 +39,9 @@ impl AuditSink for StderrSink {
     async fn emit(&self, rec: &AuditRecord) {
         match serde_json::to_string(rec) {
             Ok(js) => {
-                tracing::info!(target="dynamo_llm::audit", log_type="audit", record=%js, "audit")
+                if let Err(e) = writeln!(std::io::stderr(), "{js}") {
+                    tracing::warn!("audit: stderr write failed: {e}");
+                }
             }
             Err(e) => tracing::warn!("audit: serialize failed: {e}"),
         }
@@ -236,6 +240,7 @@ pub async fn spawn_workers_from_env(shutdown: CancellationToken) -> anyhow::Resu
                                 ) => break,
                             }
                         }
+                        sink.shutdown().await;
                         return;
                     }
                     msg = rx.recv() => {
@@ -251,6 +256,7 @@ pub async fn spawn_workers_from_env(shutdown: CancellationToken) -> anyhow::Resu
                     }
                 }
             }
+            sink.shutdown().await;
         });
     }
     tracing::info!(sinks = sink_count, "Audit sinks ready");