perf(observability): queue managed observability sanitizers

crates/core/src/api/tool.rs

@@ -30,11 +30,6 @@ use uuid::Uuid;
 
 pub use nemo_relay_types::api::tool::{ToolAttributes, ToolExecutionInterceptOutcome};
 
-fn queue_sanitized_event(event: Event, subscribers: &[EventSubscriberFn]) -> bool {
-    let scope_stack = current_scope_stack();
-    queue_sanitized_event_with_scope_stack(event, subscribers, scope_stack)
-}
-
 fn queue_sanitized_event_with_scope_stack(
     event: Event,
     subscribers: &[EventSubscriberFn],
@@ -340,8 +335,8 @@ async fn tool_call_with_subscriber_snapshot(
 ) -> Result<(ToolHandle, Vec<EventSubscriberFn>)> {
     ensure_runtime_owner()?;
     let parent_uuid = resolve_parent_uuid(params.parent);
+    let scope_stack = current_scope_stack();
     let (entries, subscribers) = {
-        let scope_stack = current_scope_stack();
         let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
         let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
             &registries.tool_sanitize_request_guardrails
@@ -356,12 +351,7 @@ async fn tool_call_with_subscriber_snapshot(
         (entries, subscribers)
     };
     let skill_loads = resolve_skill_loads(params.name, &params.args, params.metadata.as_ref());
-    let sanitized_args = NemoRelayContextState::tool_sanitize_request_snapshot_chain(
-        params.name,
-        params.args,
-        &entries,
-    )
-    .await;
+    let raw_args = params.args;
     let (handle, event, marks) = {
         let context = global_context();
         let state = context
@@ -377,7 +367,7 @@ async fn tool_call_with_subscriber_snapshot(
             .timestamp_opt(params.timestamp)
             .build();
         let handle = state.create_tool_handle(handle_params);
-        let event = state.build_tool_start_event(&handle, sanitized_args);
+        let event = state.build_tool_start_event(&handle, None);
         let marks = skill_loads
             .into_iter()
             .map(|skill_load| {
@@ -399,9 +389,29 @@ async fn tool_call_with_subscriber_snapshot(
             .collect::<Vec<_>>();
         (handle, event, marks)
     };
-    queue_sanitized_event(event, &subscribers);
+    let event_sanitizers = snapshot_event_sanitizers(&event, &scope_stack).unwrap_or_default();
+    let tool_name = handle.name.clone();
+    dispatch_transformed_event(
+        event,
+        Box::new(move |mut event| {
+            Box::pin(async move {
+                let sanitized = NemoRelayContextState::tool_sanitize_request_snapshot_chain(
+                    &tool_name, raw_args, &entries,
+                )
+                .await;
+                let mut fields = event.sanitize_fields();
+                fields.data = sanitized;
+                event.apply_sanitize_fields(fields);
+                event
+            })
+        }),
+        event_sanitizers,
+        &subscribers,
+        scope_stack.clone(),
+    );
     for mark in marks {
-        queue_sanitized_event(mark, &subscribers);
+        let sanitizers = snapshot_event_sanitizers(&mark, &scope_stack).unwrap_or_default();
+        dispatch_sanitized_event(mark, sanitizers, &subscribers, scope_stack.clone());
     }
     Ok((handle, subscribers))
 }
@@ -507,8 +517,8 @@ async fn tool_call_end_with_pending_marks(
     lifecycle_subscribers: Option<&[EventSubscriberFn]>,
 ) -> Result<()> {
     ensure_runtime_owner()?;
+    let scope_stack = current_scope_stack();
     let (entries, subscribers) = {
-        let scope_stack = current_scope_stack();
         let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
         let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
             &registries.tool_sanitize_response_guardrails
@@ -526,19 +536,6 @@ async fn tool_call_end_with_pending_marks(
         (entries, subscribers)
     };
     let subscribers = lifecycle_subscribers.unwrap_or(&subscribers);
-    let sanitized_result = NemoRelayContextState::tool_sanitize_response_snapshot_chain(
-        &params.handle.name,
-        params.result,
-        &entries,
-    )
-    .await;
-    let data = sanitized_result.and_then(|value| {
-        if value.is_null() {
-            params.data
-        } else {
-            Some(value)
-        }
-    });
     let event = {
         let context = global_context();
         let state = context
@@ -547,7 +544,7 @@ async fn tool_call_end_with_pending_marks(
         state.build_tool_end_event(
             EndToolHandleParams::builder()
                 .handle(params.handle)
-                .data_opt(data)
+                .data(Json::Null)
                 .metadata_opt(params.metadata)
                 .timestamp_opt(params.timestamp)
                 .build(),
@@ -572,9 +569,37 @@ async fn tool_call_end_with_pending_marks(
             ))
         })
         .collect::<Vec<_>>();
-    queue_sanitized_event(event, subscribers);
+    let event_sanitizers = snapshot_event_sanitizers(&event, &scope_stack).unwrap_or_default();
+    let tool_name = params.handle.name.clone();
+    let result = params.result;
+    let fallback = params.data;
+    dispatch_transformed_event(
+        event,
+        Box::new(move |mut event| {
+            Box::pin(async move {
+                let sanitized = NemoRelayContextState::tool_sanitize_response_snapshot_chain(
+                    &tool_name, result, &entries,
+                )
+                .await;
+                let mut fields = event.sanitize_fields();
+                fields.data = sanitized.and_then(|value| {
+                    if value.is_null() {
+                        fallback
+                    } else {
+                        Some(value)
+                    }
+                });
+                event.apply_sanitize_fields(fields);
+                event
+            })
+        }),
+        event_sanitizers,
+        subscribers,
+        scope_stack.clone(),
+    );
     for mark in marks {
-        queue_sanitized_event(mark, subscribers);
+        let sanitizers = snapshot_event_sanitizers(&mark, &scope_stack).unwrap_or_default();
+        dispatch_sanitized_event(mark, sanitizers, subscribers, scope_stack.clone());
     }
     Ok(())
 }

crates/core/src/stream.rs

@@ -13,28 +13,30 @@
 //! ```text
 //! raw chunk (Json) -> collector(chunk) -> Ok(()) -> yield chunk
 //!                                      -> Err(e) -> terminate stream with error
-//! upstream error -> terminate stream with error -> finalizer() -> Json -> SanitizeResponseGuardrails -> END event
-//! stream ends -> finalizer() -> Json -> SanitizeResponseGuardrails -> END event
+//! upstream error -> terminate stream with error -> finalizer() -> queue END sanitization
+//! stream ends -> finalizer() -> queue END sanitization
 //! ```
 //!
 //! The **collector** receives each chunk (Json) and can accumulate state
 //! (e.g., concatenating tokens). If the collector returns `Err`, the stream
 //! terminates immediately with that error. Upstream stream errors also
 //! terminate the stream immediately. The **finalizer** is called once when the
 //! stream terminates and returns the aggregated response as [`Json`]. That
-//! aggregated response then flows through sanitize response guardrails before
-//! being included in the END event.
+//! aggregated response is queued for sanitize response guardrails before being
+//! included in the END event. Stream termination does not await that queued
+//! observability work.
 
 use std::future::Future;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
+use chrono::Utc;
 use tokio_stream::Stream;
 
 use crate::api::event::{BaseEvent, MarkEvent};
-use crate::api::llm::LlmHandle;
 use crate::api::llm::emit_reserved_optimization_marks;
+use crate::api::llm::{EndLlmHandleParams, LlmHandle};
 use crate::api::optimization::finalize_optimization_summary;
 use crate::api::registry::Guardrail;
 use crate::api::runtime::NemoRelayContextState;
@@ -61,9 +63,9 @@ use serde_json::Map;
 /// 1. Passes each chunk to the user-supplied **collector** closure.
 ///    If the collector returns `Err`, the stream terminates with that error.
 /// 2. On stream exhaustion or explicit close, calls the **finalizer** to
-///    produce an aggregated [`Json`] response, runs sanitize response
-///    guardrails on it, then emits the LLM END event. Explicit close marks the
-///    end event as interrupted and waits for producer cleanup.
+///    produce an aggregated [`Json`] response, then queues sanitize response
+///    guardrails and LLM END event publication. Explicit close marks the end
+///    event as interrupted and waits for producer cleanup.
 ///
 /// This type is returned by [`crate::api::llm::llm_stream_call_execute`] and
 /// is usually consumed as an ordinary async stream. Consumers that stop early
@@ -84,7 +86,6 @@ pub struct LlmStreamWrapper {
     chunk_index: u64,
     ended: bool,
     close_result: Option<Result<()>>,
-    finalization: Option<tokio::task::JoinHandle<()>>,
     terminal_result: Option<Result<Json>>,
 }
 
@@ -178,7 +179,6 @@ impl LlmStreamWrapper {
             chunk_index: 0,
             ended: false,
             close_result: None,
-            finalization: None,
             terminal_result: None,
         }
     }
@@ -195,7 +195,7 @@ impl LlmStreamWrapper {
         &self.scope_stack
     }
 
-    fn finish(&mut self, background_thread: bool) {
+    fn finish(&mut self) {
         if self.ended {
             return;
         }
@@ -211,8 +211,7 @@ impl LlmStreamWrapper {
         self.handle
             .optimization_recorder
             .close_for_finalization(None);
-        self.finalization =
-            self.emit_end_event(metadata, StreamTermination::Dropped, background_thread);
+        self.emit_end_event(metadata, StreamTermination::Dropped);
     }
 
     fn finish_cleanly(&mut self) {
@@ -221,8 +220,11 @@ impl LlmStreamWrapper {
         }
         self.ended = true;
         self.inner.terminalize();
+        self.handle
+            .optimization_recorder
+            .close_for_finalization(None);
         let metadata = metadata_with_otel_status(self.metadata.clone(), "OK", None);
-        self.finalization = self.emit_end_event(metadata, StreamTermination::Complete, false);
+        self.emit_end_event(metadata, StreamTermination::Complete);
     }
 
     fn finish_with_error(&mut self, error: &FlowError) {
@@ -231,24 +233,22 @@ impl LlmStreamWrapper {
         }
         self.ended = true;
         self.inner.terminalize();
+        self.handle
+            .optimization_recorder
+            .close_for_finalization(None);
         let metadata = metadata_with_otel_error(self.metadata.clone(), error);
-        self.finalization = self.emit_end_event(metadata, StreamTermination::Failed, false);
+        self.emit_end_event(metadata, StreamTermination::Failed);
     }
 
     /// Emit the LLM END event with aggregated response data.
     ///
-    /// Calls the finalizer to produce the aggregated response, runs sanitize
-    /// response guardrails, and emits the END event.
-    fn emit_end_event(
-        &mut self,
-        metadata: Option<Json>,
-        termination: StreamTermination,
-        background_thread: bool,
-    ) -> Option<tokio::task::JoinHandle<()>> {
+    /// Calls the finalizer and queues response sanitization and END publication.
+    fn emit_end_event(&mut self, metadata: Option<Json>, termination: StreamTermination) {
         // The finalizer below runs on the caller's Tokio runtime. Register a
         // dispatcher barrier before spawning it so a synchronous subscriber
         // flush after this stream is dropped cannot overtake the END event.
         let publication_barrier = subscriber_dispatcher::register_async_publication();
+        let timestamp = Utc::now();
         let aggregated = match self.finalizer.take() {
             Some(finalizer) => finalizer(),
             None => Json::Null,
@@ -331,9 +331,17 @@ impl LlmStreamWrapper {
                 let ctx = global_context();
                 let state = ctx.read();
                 match state {
-                    Ok(state) => {
-                        Some(state.end_llm_handle(&handle, data, metadata, annotated_response))
-                    }
+                    Ok(state) => Some(
+                        state.build_llm_end_event(
+                            EndLlmHandleParams::builder()
+                                .handle(&handle)
+                                .data_opt(data)
+                                .metadata_opt(metadata)
+                                .annotated_response_opt(annotated_response)
+                                .timestamp(timestamp)
+                                .build(),
+                        ),
+                    ),
                     Err(_) => None,
                 }
             };
@@ -354,22 +362,11 @@ impl LlmStreamWrapper {
             publication_context,
             subscriber_dispatcher::with_async_publication_context(publication_barrier, finalize),
         );
-        if background_thread {
-            // `Drop` cannot await middleware and may run while the caller's
-            // executor is synchronously flushing subscribers. A process-local
-            // executor polls all detached finalizers on one shared OS thread.
-            // Pending middleware therefore does not create one thread per
-            // abandoned stream.
-            let _ = subscriber_dispatcher::spawn_background_publication(finalize);
-            return None;
-        }
-        match tokio::runtime::Handle::try_current() {
-            Ok(handle) => Some(handle.spawn(finalize)),
-            Err(_) => {
-                let _ = subscriber_dispatcher::spawn_background_publication(finalize);
-                None
-            }
-        }
+        // Stream finalization is observability-only. Queue it on the shared
+        // publication executor so stream termination does not await response
+        // or event sanitizers. The registered barrier keeps subscriber flushes
+        // ordered behind this END event.
+        let _ = subscriber_dispatcher::spawn_background_publication(finalize);
     }
 
     /// Emit a compact per-chunk receipt mark before collector processing.
@@ -437,29 +434,6 @@ impl Stream for LlmStreamWrapper {
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let this = self.as_mut().get_mut();
 
-        // The END event runs async because response and event sanitizers may
-        // await. Do not expose stream termination until that work has queued
-        // the event: callers commonly flush subscribers immediately after
-        // exhausting a stream, and that flush must include its END event.
-        if let Some(finalization) = this.finalization.as_mut() {
-            return match Pin::new(finalization).poll(cx) {
-                Poll::Pending => Poll::Pending,
-                Poll::Ready(Ok(())) => {
-                    this.finalization = None;
-                    match this.terminal_result.take() {
-                        Some(result) => Poll::Ready(Some(result)),
-                        None => Poll::Ready(None),
-                    }
-                }
-                Poll::Ready(Err(error)) => {
-                    this.finalization = None;
-                    Poll::Ready(Some(Err(FlowError::Internal(format!(
-                        "stream finalization task failed: {error}"
-                    )))))
-                }
-            };
-        }
-
         if this.ended {
             return match this.terminal_result.take() {
                 Some(result) => Poll::Ready(Some(result)),
@@ -505,12 +479,7 @@ impl LlmStreamInner for LlmStreamWrapper {
                 return result.clone();
             }
             let result = this.inner.close().await;
-            this.finish(false);
-            if let Some(finalization) = this.finalization.take() {
-                finalization.await.map_err(|error| {
-                    FlowError::Internal(format!("stream finalization task failed: {error}"))
-                })?;
-            }
+            this.finish();
             this.close_result = Some(result.clone());
             this.close_result
                 .as_ref()
@@ -788,7 +757,7 @@ fn non_empty_object(object: Map<String, Json>) -> Option<Json> {
 
 impl Drop for LlmStreamWrapper {
     fn drop(&mut self) {
-        self.finish(true);
+        self.finish();
     }
 }