From c613f260db30d1bec8965d849572fc18508ea7db Mon Sep 17 00:00:00 2001 From: Will Killian Date: Mon, 10 Aug 2026 11:46:58 -0400 Subject: [PATCH 1/5] feat: queue managed observability sanitizers Signed-off-by: Will Killian --- crates/core/src/api/llm.rs | 273 +++++++++++++----- crates/core/src/api/tool.rs | 87 ++++-- crates/core/src/stream.rs | 63 ++-- .../tests/integration/middleware_tests.rs | 267 +++++++++++++++-- docs/about-nemo-relay/architecture.mdx | 11 +- docs/about-nemo-relay/concepts/middleware.mdx | 31 +- 6 files changed, 566 insertions(+), 166 deletions(-) diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index e5b5addda..236246e24 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -22,8 +22,8 @@ use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; use crate::api::runtime::state::contextualize_stream; use crate::api::runtime::subscriber_dispatcher::{ - PendingPublication, dispatch_reserved_sanitized_event, dispatch_sanitized_event, - dispatch_transformed_event, register_pending_publication, + EventTransformFn, PendingPublication, dispatch_reserved_sanitized_event, + dispatch_sanitized_event, dispatch_transformed_event, register_pending_publication, }; use crate::api::runtime::{ EventSubscriberFn, LlmCollectorFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, @@ -420,6 +420,7 @@ fn limit_annotated_request_history_to_current_user_turn( ) } +#[cfg(test)] async fn emit_llm_start_with_subscribers( handle: &LlmHandle, request: &LlmRequest, @@ -490,6 +491,87 @@ async fn emit_llm_start_with_subscribers( Ok(()) } +fn queue_llm_start_with_subscribers( + handle: &LlmHandle, + request: &LlmRequest, + annotated_request: Option>, + request_codec: Option>, + subscribers: &[EventSubscriberFn], +) -> Result<()> { + ensure_runtime_owner()?; + let scope_stack = handle.captured_scope_stack().clone(); + let (entries, agent_is_fresh, full_payloads_enabled) = { + let mut scope_guard = scope_stack.write().expect("scope stack lock poisoned"); + let scope_locals = scope_guard.collect_scope_local_registries(|registries| { + ®istries.llm_sanitize_request_guardrails + }); + let context = global_context(); + let state = context + .read() + .map_err(|error| FlowError::Internal(error.to_string()))?; + let entries = state.llm_sanitize_request_entries(&scope_locals); + let full_payloads_enabled = state.observability_full_payloads_enabled; + drop(state); + let agent_is_fresh = scope_guard.take_agent_freshness(handle.parent_uuid); + (entries, agent_is_fresh, full_payloads_enabled) + }; + let observable_request = remove_observability_credential_headers(request.clone()); + let queued_handle = handle.clone(); + let event = { + let context = global_context(); + let state = context + .read() + .map_err(|error| FlowError::Internal(error.to_string()))?; + state.build_llm_start_event(handle, None, None) + }; + let event_sanitizers = snapshot_event_sanitizers(&event, &scope_stack).unwrap_or_default(); + dispatch_transformed_event( + event, + Box::new(move |event| { + Box::pin(async move { + let mut sanitized_request = + NemoRelayContextState::llm_sanitize_request_snapshot_chain( + observable_request.clone(), + LlmSanitizeRequestContext::for_request_codec(request_codec.clone()), + &entries, + ) + .await; + let request_changed = sanitized_request + .as_ref() + .is_some_and(|sanitized| sanitized != &observable_request); + let mut annotation = match (sanitized_request.as_ref(), request_codec.as_deref()) { + (Some(sanitized), Some(codec)) if request_changed => { + codec.decode(sanitized).ok().map(Arc::new) + } + (Some(_), _) if !request_changed => annotated_request, + _ => None, + }; + if !full_payloads_enabled + && !agent_is_fresh + && let Some(sanitized_request) = sanitized_request.as_mut() + { + project_llm_request_to_current_user_turn( + sanitized_request, + &mut annotation, + request_codec.as_deref(), + ); + } + let input = sanitized_request + .as_ref() + .and_then(|request| serde_json::to_value(request).ok()); + global_context() + .read() + .map(|state| state.build_llm_start_event(&queued_handle, input, annotation)) + .unwrap_or(event) + }) + }), + event_sanitizers, + subscribers, + scope_stack, + ); + Ok(()) +} + fn remove_observability_credential_headers(mut request: LlmRequest) -> LlmRequest { request.headers.retain(|name, _| { !OBSERVABILITY_CREDENTIAL_HEADERS @@ -846,7 +928,6 @@ pub fn llm_call(params: LlmCallParams<'_>) -> Result { #[derive(Clone, Copy)] struct LlmCallEndBehavior { - response_codec_errors_fatal: bool, attach_estimated_cost: bool, } @@ -856,6 +937,18 @@ struct LlmEndPayload { decode_error: Option, } +/// Queue a provisional LLM END event and replace its observability-only +/// payload on the serial publication path before event sanitizers run. +fn queue_llm_end_event( + event: Event, + transform: EventTransformFn, + subscribers: &[EventSubscriberFn], + scope_stack: ScopeStackHandle, +) { + let event_sanitizers = snapshot_event_sanitizers(&event, &scope_stack).unwrap_or_default(); + dispatch_transformed_event(event, transform, event_sanitizers, subscribers, scope_stack); +} + async fn build_llm_end_payload( handle: &LlmHandle, response: Json, @@ -1012,7 +1105,6 @@ pub fn llm_call_end(params: LlmCallEndParams<'_>) -> Result<()> { response_codec, &entries, LlmCallEndBehavior { - response_codec_errors_fatal: false, attach_estimated_cost: false, }, ) @@ -1081,41 +1173,65 @@ async fn llm_call_end_with_behavior( (entries, subscribers) }; handle.optimization_recorder.close_for_finalization(None); - emit_optimization_marks(handle, &subscribers).await; - let payload = build_llm_end_payload( - handle, - response, - data, - annotated_response, - response_codec, - &entries, - behavior, - ) - .await; + enqueue_optimization_marks(handle, &subscribers); + let queued_handle = handle.clone(); let event = { let context = global_context(); let state = context .read() .map_err(|error| FlowError::Internal(error.to_string()))?; - let end_metadata = metadata_with_otel_status(metadata, "OK", None); + let end_metadata = metadata_with_otel_status(metadata.clone(), "OK", None); state.build_llm_end_event( EndLlmHandleParams::builder() .handle(handle) - .data_opt(payload.data) + .data(Json::Null) .metadata_opt(end_metadata) - .annotated_response_opt(payload.annotated_response) .timestamp_opt(timestamp) .build(), ) }; - queue_sanitized_event_with_scope_stack(event, &subscribers, handle.captured_scope_stack()); - if let Some(error) = payload.decode_error - && behavior.response_codec_errors_fatal - { - Err(error) - } else { - Ok(()) - } + let scope_stack = handle.captured_scope_stack().clone(); + queue_llm_end_event( + event, + Box::new(move |event| { + Box::pin(async move { + let payload = build_llm_end_payload( + &queued_handle, + response, + data, + annotated_response, + response_codec, + &entries, + behavior, + ) + .await; + if let Some(error) = payload.decode_error { + log::error!( + target: "nemo_relay.runtime", + event = "managed_llm_response_codec_failed"; + "Managed LLM response annotation failed during queued publication: {error}" + ); + } + let context = global_context(); + let Ok(state) = context.read() else { + return event; + }; + let end_metadata = metadata_with_otel_status(metadata, "OK", None); + state.build_llm_end_event( + EndLlmHandleParams::builder() + .handle(&queued_handle) + .data_opt(payload.data) + .metadata_opt(end_metadata) + .annotated_response_opt(payload.annotated_response) + .timestamp_opt(timestamp) + .build(), + ) + }) + }), + &subscribers, + scope_stack, + ); + Ok(()) } #[cfg(test)] @@ -1183,46 +1299,70 @@ async fn emit_llm_end_without_output( let entries = state.llm_sanitize_response_entries(&scope_locals); (entries, subscribers) }; - let had_fallback_data = handle.data.is_some(); - let data = if let Some(data) = handle.data.clone() { - NemoRelayContextState::llm_sanitize_response_snapshot_chain( - data, - LlmSanitizeResponseContext::for_response_codec(response_codec), - &entries, - ) - .await - } else { - None - }; - let annotation_omitted = - (had_fallback_data && data.is_none()) || data.as_ref().is_some_and(Json::is_null); handle.optimization_recorder.close_for_finalization(None); - emit_optimization_marks(handle, &subscribers).await; - let pricing = crate::codec::response::active_pricing_resolver(); - let annotated_response = (!annotation_omitted) - .then(|| { - finalize_optimization_summary( - &handle.optimization_recorder, - None, - handle.model_name.as_deref(), - &pricing, - ) - }) - .flatten() - .map(|summary| { - Arc::new(AnnotatedLlmResponse { - optimization_summary: Some(summary), - ..AnnotatedLlmResponse::default() - }) - }); + enqueue_optimization_marks(handle, &subscribers); + let queued_handle = handle.clone(); + let fallback_data = handle.data.clone(); let event = { let context = global_context(); let state = context .read() .map_err(|error| FlowError::Internal(error.to_string()))?; - state.end_llm_handle(handle, data, metadata, annotated_response) + state.build_llm_end_event( + EndLlmHandleParams::builder() + .handle(handle) + .data(Json::Null) + .metadata_opt(metadata.clone()) + .build(), + ) }; - queue_sanitized_event_with_scope_stack(event, &subscribers, handle.captured_scope_stack()); + let scope_stack = handle.captured_scope_stack().clone(); + queue_llm_end_event( + event, + Box::new(move |event| { + Box::pin(async move { + let had_fallback_data = fallback_data.is_some(); + let data = match fallback_data { + Some(data) => { + NemoRelayContextState::llm_sanitize_response_snapshot_chain( + data, + LlmSanitizeResponseContext::for_response_codec(response_codec), + &entries, + ) + .await + } + None => None, + }; + let annotation_omitted = (had_fallback_data && data.is_none()) + || data.as_ref().is_some_and(Json::is_null); + let pricing = crate::codec::response::active_pricing_resolver(); + let annotated_response = (!annotation_omitted) + .then(|| { + finalize_optimization_summary( + &queued_handle.optimization_recorder, + None, + queued_handle.model_name.as_deref(), + &pricing, + ) + }) + .flatten() + .map(|summary| { + Arc::new(AnnotatedLlmResponse { + optimization_summary: Some(summary), + ..AnnotatedLlmResponse::default() + }) + }); + global_context() + .read() + .map(|state| { + state.end_llm_handle(&queued_handle, data, metadata, annotated_response) + }) + .unwrap_or(event) + }) + }), + &subscribers, + scope_stack, + ); Ok(()) } @@ -1477,15 +1617,14 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); snapshot_event_subscribers(scope_guard.collect_scope_local_subscribers())? }; - emit_llm_start_with_subscribers( + inject_traceparent(&mut intercepted_request, handle.uuid)?; + queue_llm_start_with_subscribers( &handle, &intercepted_request, annotated_request.clone(), request_codec.clone(), &lifecycle_subscribers, - ) - .await?; - inject_traceparent(&mut intercepted_request, handle.uuid)?; + )?; emit_pending_request_marks(&handle, pending_marks, &lifecycle_subscribers).await?; handle .optimization_recorder @@ -1531,7 +1670,6 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { .response_codec_opt(response_codec) .build(), LlmCallEndBehavior { - response_codec_errors_fatal: false, attach_estimated_cost: true, }, Some(&lifecycle_subscribers), @@ -1689,15 +1827,14 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); snapshot_event_subscribers(scope_guard.collect_scope_local_subscribers())? }; - emit_llm_start_with_subscribers( + inject_traceparent(&mut intercepted_request, handle.uuid)?; + queue_llm_start_with_subscribers( &handle, &intercepted_request, annotated_request, request_codec.clone(), &lifecycle_subscribers, - ) - .await?; - inject_traceparent(&mut intercepted_request, handle.uuid)?; + )?; emit_pending_request_marks(&handle, pending_marks, &lifecycle_subscribers).await?; handle .optimization_recorder diff --git a/crates/core/src/api/tool.rs b/crates/core/src/api/tool.rs index 61bd6f922..a84af3327 100644 --- a/crates/core/src/api/tool.rs +++ b/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], @@ -356,12 +351,7 @@ async fn tool_call_with_subscriber_snapshot( (entries, subscribers) }; let skill_loads = resolve_skill_loads(params.name, ¶ms.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,30 @@ async fn tool_call_with_subscriber_snapshot( .collect::>(); (handle, event, marks) }; - queue_sanitized_event(event, &subscribers); + let scope_stack = current_scope_stack(); + 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)) } @@ -526,19 +537,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( - ¶ms.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 +545,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 +570,38 @@ async fn tool_call_end_with_pending_marks( )) }) .collect::>(); - queue_sanitized_event(event, subscribers); + let scope_stack = current_scope_stack(); + 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(()) } diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index d2defea3e..bed67944f 100644 --- a/crates/core/src/stream.rs +++ b/crates/core/src/stream.rs @@ -13,8 +13,8 @@ //! ```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 @@ -22,19 +22,21 @@ //! 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 @@ -243,12 +245,13 @@ impl LlmStreamWrapper { &mut self, metadata: Option, termination: StreamTermination, - background_thread: bool, + _background_thread: bool, ) -> Option> { // 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 +334,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 +365,12 @@ 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); + None } /// Emit a compact per-chunk receipt mark before collector processing. @@ -437,10 +438,8 @@ impl Stream for LlmStreamWrapper { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 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. + // Retain support for an already-scheduled finalization task while the + // stream is being polled. if let Some(finalization) = this.finalization.as_mut() { return match Pin::new(finalization).poll(cx) { Poll::Pending => Poll::Pending, diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index 0fba0173b..555e8cd37 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -16,6 +16,7 @@ use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; +use chrono::Utc; mod test_support; use test_support::{ready, ready_result}; @@ -3519,6 +3520,7 @@ async fn test_tool_middleware_callbacks_run_without_registry_or_scope_locks() { .await .unwrap(); assert_eq!(result["ok"], true); + flush_subscribers().unwrap(); assert_middleware_callback_labels( &callbacks, &[ @@ -3799,6 +3801,7 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { chunk.unwrap(); } stream.close().await.unwrap(); + flush_subscribers().unwrap(); assert_middleware_callback_labels( &callbacks, &[ @@ -3941,25 +3944,19 @@ async fn test_full_pipeline_integration() { .await .unwrap(); - // Verify the pipeline order: - // 1. conditional (runs on raw args, before intercepts) - // 2. request_intercept (transforms args) - // 3. sanitize_request (inside tool_call) - // 4. execution_intercept -> original_execution - // 5. sanitize_response (inside tool_call_end) + flush_subscribers().unwrap(); + + // Application middleware remains ordered on the managed path. Payload + // sanitizers run on the publication path, where request still precedes + // response but may race with tool execution. let recorded = order.lock().unwrap(); - assert_eq!( - *recorded, - vec![ - "conditional", - "request_intercept", - "sanitize_request", - "execution_intercept", - "original_execution", - "sanitize_response", - ], - "Full pipeline should execute in the correct order" - ); + let index = |name: &str| recorded.iter().position(|entry| entry == name).unwrap(); + assert!(index("conditional") < index("request_intercept")); + assert!(index("request_intercept") < index("execution_intercept")); + assert!(index("execution_intercept") < index("original_execution")); + assert!(index("request_intercept") < index("sanitize_request")); + assert!(index("sanitize_request") < index("sanitize_response")); + assert!(index("original_execution") < index("sanitize_response")); // Verify the request intercept's modification persists through the pipeline assert_eq!(result["intercepted"], true); @@ -4419,6 +4416,240 @@ async fn test_managed_llm_event_sanitizers_run_off_execution_path_in_fifo_order( deregister_subscriber("managed_async_publication_observer").unwrap(); } +#[tokio::test] +async fn test_managed_llm_payload_sanitizers_are_queued_off_execution_path() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let request_started = Arc::new(tokio::sync::Notify::new()); + let request_release = Arc::new(tokio::sync::Notify::new()); + register_llm_sanitize_request_guardrail( + "managed_queued_llm_request", + 1, + Arc::new({ + let request_started = Arc::clone(&request_started); + let request_release = Arc::clone(&request_release); + move |request, _context| { + let request_started = Arc::clone(&request_started); + let request_release = Arc::clone(&request_release); + Box::pin(async move { + request_started.notify_one(); + request_release.notified().await; + Ok(Some(request)) + }) + } + }), + ) + .unwrap(); + let response_started = Arc::new(tokio::sync::Notify::new()); + let response_release = Arc::new(tokio::sync::Notify::new()); + register_llm_sanitize_response_guardrail( + "managed_queued_llm_response", + 1, + Arc::new({ + let response_started = Arc::clone(&response_started); + let response_release = Arc::clone(&response_release); + move |response, _context| { + let response_started = Arc::clone(&response_started); + let response_release = Arc::clone(&response_release); + Box::pin(async move { + response_started.notify_one(); + response_release.notified().await; + Ok(Some(response)) + }) + } + }), + ) + .unwrap(); + register_subscriber("managed_queued_llm_observer", Arc::new(|_| {})).unwrap(); + + let call = tokio::spawn(async { + llm_call_execute( + LlmCallExecuteParams::builder() + .name("managed-queued-llm") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({"prompt": "hello"}), + }) + .func(Arc::new(|_| { + Box::pin(async { Ok(json!({"response": "done"})) }) + })) + .build(), + ) + .await + }); + tokio::time::timeout( + std::time::Duration::from_secs(2), + request_started.notified(), + ) + .await + .expect("managed request sanitizer did not start"); + let result = tokio::time::timeout(std::time::Duration::from_secs(1), call) + .await + .expect("request sanitizer blocked managed provider execution") + .expect("managed call task should join") + .expect("managed call should succeed"); + assert_eq!(result, json!({"response": "done"})); + + request_release.notify_one(); + tokio::time::timeout( + std::time::Duration::from_secs(2), + response_started.notified(), + ) + .await + .expect("managed response sanitizer did not start"); + assert_flush_waits_for_pending_completion(|| response_release.notify_one()); + deregister_llm_sanitize_request_guardrail("managed_queued_llm_request").unwrap(); + deregister_llm_sanitize_response_guardrail("managed_queued_llm_response").unwrap(); + deregister_subscriber("managed_queued_llm_observer").unwrap(); +} + +#[tokio::test] +async fn test_managed_tool_payload_sanitizers_are_queued_off_execution_path() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let sanitizer_started = Arc::new(tokio::sync::Notify::new()); + let sanitizer_release = Arc::new(tokio::sync::Notify::new()); + register_tool_sanitize_request_guardrail( + "managed_queued_tool_request", + 1, + Arc::new({ + let sanitizer_started = Arc::clone(&sanitizer_started); + let sanitizer_release = Arc::clone(&sanitizer_release); + move |_name, args| { + let sanitizer_started = Arc::clone(&sanitizer_started); + let sanitizer_release = Arc::clone(&sanitizer_release); + Box::pin(async move { + sanitizer_started.notify_one(); + sanitizer_release.notified().await; + Ok(args) + }) + } + }), + ) + .unwrap(); + register_subscriber("managed_queued_tool_observer", Arc::new(|_| {})).unwrap(); + + let call = tokio::spawn(async { + tool_call_execute( + nemo_relay::api::tool::ToolCallExecuteParams::builder() + .name("managed-queued-tool") + .args(json!({"input": true})) + .func(Arc::new(|args| Box::pin(async move { Ok(args) }))) + .build(), + ) + .await + }); + tokio::time::timeout( + std::time::Duration::from_secs(2), + sanitizer_started.notified(), + ) + .await + .expect("managed tool request sanitizer did not start"); + let result = tokio::time::timeout(std::time::Duration::from_secs(1), call) + .await + .expect("request sanitizer blocked managed tool execution") + .expect("managed tool task should join") + .expect("managed tool call should succeed"); + assert_eq!(result, json!({"input": true})); + + sanitizer_release.notify_one(); + flush_subscribers().unwrap(); + deregister_tool_sanitize_request_guardrail("managed_queued_tool_request").unwrap(); + deregister_subscriber("managed_queued_tool_observer").unwrap(); +} + +#[tokio::test] +async fn test_stream_termination_does_not_await_response_sanitization() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let sanitizer_started = Arc::new(tokio::sync::Notify::new()); + let sanitizer_release = Arc::new(tokio::sync::Notify::new()); + let events = Arc::new(Mutex::new(Vec::::new())); + register_llm_sanitize_response_guardrail( + "managed_queued_stream_response", + 1, + Arc::new({ + let sanitizer_started = Arc::clone(&sanitizer_started); + let sanitizer_release = Arc::clone(&sanitizer_release); + move |response, _context| { + let sanitizer_started = Arc::clone(&sanitizer_started); + let sanitizer_release = Arc::clone(&sanitizer_release); + Box::pin(async move { + sanitizer_started.notify_one(); + sanitizer_release.notified().await; + Ok(Some(response)) + }) + } + }), + ) + .unwrap(); + register_subscriber( + "managed_queued_stream_observer", + Arc::new({ + let events = Arc::clone(&events); + move |event| events.lock().unwrap().push(event.clone()) + }), + ) + .unwrap(); + + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("managed-queued-stream") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({"prompt": "hello"}), + }) + .func(Arc::new(|_| { + Box::pin(async { + Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ + "chunk": "done" + }))]))) + }) + })) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| json!({"response": "done"}))) + .build(), + ) + .await + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while let Some(item) = stream.next().await { + item.unwrap(); + } + }) + .await + .expect("response sanitizer blocked stream termination"); + tokio::time::timeout( + std::time::Duration::from_secs(2), + sanitizer_started.notified(), + ) + .await + .expect("stream response sanitizer did not start"); + let terminal_timestamp = Utc::now(); + + assert_flush_waits_for_pending_completion(|| sanitizer_release.notify_one()); + let end = events + .lock() + .unwrap() + .iter() + .find(|event| { + event.name() == "managed-queued-stream" + && event.scope_category() == Some(ScopeCategory::End) + }) + .cloned() + .expect("stream END event should be published after flush"); + assert!(*end.timestamp() <= terminal_timestamp); + + deregister_llm_sanitize_response_guardrail("managed_queued_stream_response").unwrap(); + deregister_subscriber("managed_queued_stream_observer").unwrap(); +} + #[tokio::test] async fn test_stream_response_sanitizer_nested_mark_precedes_end_event() { let _lock = TEST_MUTEX.lock().unwrap(); diff --git a/docs/about-nemo-relay/architecture.mdx b/docs/about-nemo-relay/architecture.mdx index d11e8bf64..03c28dd51 100644 --- a/docs/about-nemo-relay/architecture.mdx +++ b/docs/about-nemo-relay/architecture.mdx @@ -115,10 +115,13 @@ model: - [**Subscribers and exporters**](/about-nemo-relay/concepts/subscribers#common-subscriber-roles) consume those snapshots. -Every managed tool or LLM call resolves the middleware visible from the active -scope before it executes. When the runtime emits an event, it records the active -scope UUID as parentage. The scope stack changes as work opens and closes; the -parent-linked event records remain available to subscribers. +Every managed tool or LLM call resolves the conditional and intercept middleware +visible from the active scope before it executes. Payload sanitization and event +publication are queued observability work: they do not delay the real callback +or result, and `flush_subscribers()` is the explicit delivery barrier. When the +runtime emits an event, it records the active scope UUID as parentage. The scope +stack changes as work opens and closes; the parent-linked event records remain +available to subscribers. ## Main Runtime Pieces diff --git a/docs/about-nemo-relay/concepts/middleware.mdx b/docs/about-nemo-relay/concepts/middleware.mdx index b06c58ab1..82a7e9b5f 100644 --- a/docs/about-nemo-relay/concepts/middleware.mdx +++ b/docs/about-nemo-relay/concepts/middleware.mdx @@ -25,9 +25,10 @@ All middleware families are asynchronous in the Rust runtime. Rust callbacks return a future, and Node callbacks may return a value or a Promise. Python registrations accept callbacks that return a value or an awaitable when invoked through an asynchronous Relay API or queued event publication. Worker and -native-plugin middleware can also complete asynchronously. Relay awaits entries -sequentially in priority order, so later callbacks observe earlier middleware -output. +native-plugin middleware can also complete asynchronously. Within each +middleware chain, Relay awaits entries sequentially in priority order so later +callbacks observe earlier middleware output. Payload and event sanitizer chains +run on the queued publication path and do not delay managed execution. The experimental raw C FFI and Go binding retain synchronous middleware callbacks. Relay invokes each callback on a native thread and waits for it to @@ -39,12 +40,13 @@ Synchronous standalone Python calls cannot drive an awaitable callback. Call the same standalone helper from a running event loop and await the returned value instead. -Managed execution is asynchronous because its result depends on middleware -completion. Python standalone conditional and request-intercept helpers return -a direct value outside an event loop and an awaitable inside one. Manual lifecycle -APIs (`tool_call`, `tool_call_end`, `llm_call`, and `llm_call_end`) remain -synchronous: they create or close their handle immediately and queue observability -work rather than awaiting it. +Managed execution is asynchronous because its result depends on conditional +guardrails and intercept completion. Python standalone conditional and +request-intercept helpers return a direct value outside an event loop and an +awaitable inside one. Managed and manual lifecycle APIs queue observability +sanitization and publication rather than awaiting it. Manual lifecycle APIs +(`tool_call`, `tool_call_end`, `llm_call`, and `llm_call_end`) remain +synchronous and create or close their handle immediately. Event sanitizers, conditional-execution guardrails, request intercepts, @@ -217,11 +219,12 @@ arguments passed to the callback or the real value returned to the caller. ## Queued Event Publication -Scope operations, marks, and manual tool/LLM lifecycle calls never become -awaitable because an event sanitizer is asynchronous. At emission time Relay -snapshots the event, visible sanitizer chain, and subscribers, then places the -work on a serial dispatcher. The dispatcher awaits sanitizers and publishes the -event later in FIFO order. +Scope operations, marks, and manual or managed tool/LLM lifecycle calls do not +await observability sanitizers. At emission time Relay snapshots the event-only +payload, visible sanitizer chains, and subscribers, then places the work on a +serial dispatcher. The dispatcher awaits the specialized tool or LLM payload +sanitizers, then the event sanitizers, and publishes the event later in FIFO +order. Subscriber and exporter delivery is therefore delayed, while start/end/mark order is preserved. Closing a scope or deregistering middleware after emission From 35692a780de5565d3e37a6ae8ba7f927744b3f32 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Mon, 10 Aug 2026 12:05:36 -0400 Subject: [PATCH 2/5] docs: clarify queued sanitizer ordering Signed-off-by: Will Killian --- docs/about-nemo-relay/concepts/middleware.mdx | 45 ++++++++++++------- 1 file changed, 29 insertions(+), 16 deletions(-) diff --git a/docs/about-nemo-relay/concepts/middleware.mdx b/docs/about-nemo-relay/concepts/middleware.mdx index 82a7e9b5f..20e0b52a3 100644 --- a/docs/about-nemo-relay/concepts/middleware.mdx +++ b/docs/about-nemo-relay/concepts/middleware.mdx @@ -243,9 +243,8 @@ flowchart LR subgraph RequestPhase[Request Phase] Conditional[Conditional Guardrails] RequestIntercepts[Request Intercepts] - RequestSanitizers[Request Sanitizers] - StartEvent[Scope-Start Sanitizers and Start Event] - Conditional --> RequestIntercepts --> RequestSanitizers --> StartEvent + QueueStart[Enqueue Start Event Snapshot] + Conditional --> RequestIntercepts --> QueueStart end subgraph ExecutionPhase[Execution Phase] @@ -255,28 +254,42 @@ flowchart LR end subgraph ResponsePhase[Response Phase] + QueueEnd[Enqueue End Event Snapshot] + Return[Return Real Result] + QueueEnd --> Return + end + + QueueStart --> ExecutionIntercepts + Callback --> QueueEnd + + subgraph PublicationPath[Serial Publication Path] + RequestSanitizers[Request Sanitizers] + StartEvent[Scope-Start Sanitizers and Deliver Start] ResponseSanitizers[Response Sanitizers] - EndEvent[Scope-End Sanitizers and End Event] - ResponseSanitizers --> EndEvent + EndEvent[Scope-End Sanitizers and Deliver End] + RequestSanitizers --> StartEvent --> ResponseSanitizers --> EndEvent end - StartEvent --> ExecutionIntercepts - Callback --> ResponseSanitizers + QueueStart -.-> RequestSanitizers + QueueEnd -.-> ResponseSanitizers ``` The phases run as follows: 1. **Request phase:** Conditional-execution guardrails allow the call, request - intercepts rewrite the real request, request sanitizers rewrite the - observability copy, and scope-start sanitizers run before Relay enqueues the - start event. + intercepts rewrite the real request, and Relay snapshots and enqueues the + start event's observability copy. 2. **Execution phase:** Execution intercepts wrap or replace the real callback. -3. **Response phase:** Response sanitizers rewrite the observability copy, and - scope-end sanitizers run before Relay enqueues the end event. - -The start event is enqueued before execution begins. The async subscriber -dispatcher delivers start and end snapshots later in FIFO order; managed -execution does not wait for subscriber or exporter callbacks. +3. **Response phase:** Relay snapshots and enqueues the end event's + observability copy, then returns the real result. +4. **Publication path:** The serial dispatcher runs request or response + sanitizers, then the matching scope-event sanitizers, and finally delivers + the event to subscribers and exporters. + +The start event is submitted before execution begins, but its sanitizers run +later on the publication path and may overlap application execution. The serial +dispatcher preserves start/end delivery order, and `flush_subscribers()` waits +for queued sanitization and delivery when a caller needs that barrier. This ordering preserves the distinction between the families: From a6ab73e4437a2ca38b72a426551dbe15054e7724 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Mon, 10 Aug 2026 12:53:50 -0400 Subject: [PATCH 3/5] test: flush queued sanitizer publications Signed-off-by: Will Killian --- .../tests/integration/middleware_tests.rs | 28 +++++++++++-------- crates/core/tests/integration/stream_tests.rs | 6 ++++ crates/core/tests/unit/llm_api_tests.rs | 2 +- crates/node/tests/llm_tests.mjs | 4 ++- go/nemo_relay/llm_test.go | 3 ++ go/nemo_relay/scope_local_test.go | 12 ++++++++ go/nemo_relay/tools_test.go | 3 ++ python/tests/test_llm.py | 1 + python/tests/test_scope_local.py | 2 ++ 9 files changed, 47 insertions(+), 14 deletions(-) diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index 555e8cd37..877cc6a0c 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -171,16 +171,20 @@ fn captured_events_snapshot(events: &Arc>>) -> Vec { } fn assert_middleware_callback_locks_are_free() { - let context = global_context(); + let scope_stack = current_scope_stack(); assert!( - context.try_write().is_ok(), - "middleware callback ran while the global registry lock was held" + scope_stack.try_write().is_ok(), + "middleware callback ran while its scope stack lock was held" ); +} +/// Queued payload sanitizers may overlap later hot-path registry reads. They +/// still must not run while holding their captured scope-stack lock. +fn assert_queued_sanitizer_scope_lock_is_free() { let scope_stack = current_scope_stack(); assert!( scope_stack.try_write().is_ok(), - "middleware callback ran while the scope stack lock was held" + "queued sanitizer ran while the scope stack lock was held" ); } @@ -3440,7 +3444,7 @@ async fn test_tool_middleware_callbacks_run_without_registry_or_scope_locks() { 1, Arc::new(move |_, args| { record_middleware_callback(&tracked, "tool_sanitize_request_global"); - assert_middleware_callback_locks_are_free(); + assert_queued_sanitizer_scope_lock_is_free(); ready(args) }), ) @@ -3452,7 +3456,7 @@ async fn test_tool_middleware_callbacks_run_without_registry_or_scope_locks() { 2, Arc::new(move |_, args| { record_middleware_callback(&tracked, "tool_sanitize_request_scope"); - assert_middleware_callback_locks_are_free(); + assert_queued_sanitizer_scope_lock_is_free(); ready(args) }), ) @@ -3486,7 +3490,7 @@ async fn test_tool_middleware_callbacks_run_without_registry_or_scope_locks() { 1, Arc::new(move |_, result| { record_middleware_callback(&tracked, "tool_sanitize_response_global"); - assert_middleware_callback_locks_are_free(); + assert_queued_sanitizer_scope_lock_is_free(); ready(result) }), ) @@ -3498,7 +3502,7 @@ async fn test_tool_middleware_callbacks_run_without_registry_or_scope_locks() { 2, Arc::new(move |_, result| { record_middleware_callback(&tracked, "tool_sanitize_response_scope"); - assert_middleware_callback_locks_are_free(); + assert_queued_sanitizer_scope_lock_is_free(); ready(result) }), ) @@ -3655,7 +3659,7 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { 1, Arc::new(move |request, _context| { record_middleware_callback(&tracked, "llm_sanitize_request_global"); - assert_middleware_callback_locks_are_free(); + assert_queued_sanitizer_scope_lock_is_free(); ready(Some(request)) }), ) @@ -3667,7 +3671,7 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { 2, Arc::new(move |request, _context| { record_middleware_callback(&tracked, "llm_sanitize_request_scope"); - assert_middleware_callback_locks_are_free(); + assert_queued_sanitizer_scope_lock_is_free(); ready(Some(request)) }), ) @@ -3724,7 +3728,7 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { 1, Arc::new(move |response, _context| { record_middleware_callback(&tracked, "llm_sanitize_response_global"); - assert_middleware_callback_locks_are_free(); + assert_queued_sanitizer_scope_lock_is_free(); ready(Some(response)) }), ) @@ -3736,7 +3740,7 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { 2, Arc::new(move |response, _context| { record_middleware_callback(&tracked, "llm_sanitize_response_scope"); - assert_middleware_callback_locks_are_free(); + assert_queued_sanitizer_scope_lock_is_free(); ready(Some(response)) }), ) diff --git a/crates/core/tests/integration/stream_tests.rs b/crates/core/tests/integration/stream_tests.rs index 0b66a0ab2..512f5543d 100644 --- a/crates/core/tests/integration/stream_tests.rs +++ b/crates/core/tests/integration/stream_tests.rs @@ -566,6 +566,7 @@ async fn stream_termination_modes_close_accounting_without_losing_evidence() { while let Some(item) = clean.next().await { item.unwrap(); } + flush_subscribers().unwrap(); assert!(!clean_recorder.record(LlmOptimizationContribution::new("late", "test"))); let (before_error_handle, before_error_recorder) = @@ -581,6 +582,7 @@ async fn stream_termination_modes_close_accounting_without_losing_evidence() { None, ); assert!(error_before.next().await.unwrap().is_err()); + flush_subscribers().unwrap(); assert!(!before_error_recorder.record(LlmOptimizationContribution::new("late", "test"))); let (after_error_handle, after_error_recorder) = @@ -600,6 +602,7 @@ async fn stream_termination_modes_close_accounting_without_losing_evidence() { ); assert!(error_after.next().await.unwrap().is_ok()); assert!(error_after.next().await.unwrap().is_err()); + flush_subscribers().unwrap(); assert!(!after_error_recorder.record(LlmOptimizationContribution::new("late", "test"))); let (drop_before_handle, drop_before_recorder) = @@ -615,6 +618,7 @@ async fn stream_termination_modes_close_accounting_without_losing_evidence() { None, ); drop(drop_before); + flush_subscribers().unwrap(); assert!(!drop_before_recorder.record(LlmOptimizationContribution::new("late", "test"))); let (drop_after_handle, drop_after_recorder) = @@ -635,8 +639,10 @@ async fn stream_termination_modes_close_accounting_without_losing_evidence() { ); assert!(drop_after.next().await.unwrap().is_ok()); drop(drop_after); + flush_subscribers().unwrap(); assert!(!drop_after_recorder.record(LlmOptimizationContribution::new("late", "test"))); + flush_subscribers().unwrap(); let events = captured_snapshot(&events); let summary_for = |name: &str| { events diff --git a/crates/core/tests/unit/llm_api_tests.rs b/crates/core/tests/unit/llm_api_tests.rs index 45b26a854..6d8ff05ad 100644 --- a/crates/core/tests/unit/llm_api_tests.rs +++ b/crates/core/tests/unit/llm_api_tests.rs @@ -1286,8 +1286,8 @@ fn projection_encode_failures_do_not_block_managed_or_streaming_calls() { } }); - assert_eq!(projection_attempts.load(Ordering::Relaxed), 2); flush_subscribers().unwrap(); + assert_eq!(projection_attempts.load(Ordering::Relaxed), 2); assert!(deregister_subscriber("projection-encode-failure").unwrap()); let events = events.lock().unwrap(); for name in [ diff --git a/crates/node/tests/llm_tests.mjs b/crates/node/tests/llm_tests.mjs index 6a2ceea12..07312de1b 100644 --- a/crates/node/tests/llm_tests.mjs +++ b/crates/node/tests/llm_tests.mjs @@ -400,9 +400,9 @@ describe('LLM guardrails', () => { try { const result = await llmCallExecute('contextual_sanitize_llm', makeNative(), () => ({ ok: true })); assert.deepEqual(result, { ok: true }); + await flushSubscribers(); assert.equal(requestContextChecked, true); assert.equal(responseContextChecked, true); - await flushSubscribers(); const start = events.find( (event) => event.name === 'contextual_sanitize_llm' && event.scope_category === 'start', ); @@ -505,6 +505,7 @@ describe('LLM guardrails', () => { ({ annotated, original }) => codec.encode(annotated, original), codec.decodeResponse.bind(codec), ); + await flushSubscribers(); assert.deepEqual(result, response); assert.equal(requestDecoded, true); assert.equal(responseDecoded, true); @@ -615,6 +616,7 @@ describe('LLM guardrails', () => { const stream = await execution; assert.deepEqual(await stream.next(), { token: 'done' }); assert.equal(await stream.next(), null); + await flushSubscribers(); assert.deepEqual(observed, [invocationScope.uuid, invocationScope.uuid]); } finally { lib.withScopeStack(invocationStack, () => lib.popScope(invocationScope)); diff --git a/go/nemo_relay/llm_test.go b/go/nemo_relay/llm_test.go index f9ab2de5d..a663808d3 100644 --- a/go/nemo_relay/llm_test.go +++ b/go/nemo_relay/llm_test.go @@ -445,6 +445,9 @@ func TestLlmSanitizersResolveDirectionalCodecs(t *testing.T) { if err != nil { t.Fatalf(llmCallExecuteFailed, err) } + if err := FlushSubscribers(); err != nil { + t.Fatalf(llmFlushSubscribersFailed, err) + } assertResolvedCodecsExpire(t, callbackState.snapshot(), response) } diff --git a/go/nemo_relay/scope_local_test.go b/go/nemo_relay/scope_local_test.go index 565d0552e..567c4abf4 100644 --- a/go/nemo_relay/scope_local_test.go +++ b/go/nemo_relay/scope_local_test.go @@ -90,6 +90,9 @@ func assertScopeLocalCallbackDeregisters( if err := runBefore(); err != nil { t.Fatalf("%s before deregister failed: %v", label, err) } + if err := FlushSubscribers(); err != nil { + t.Fatalf(scopeLocalFlushSubscribersFailed, err) + } if *calls != 1 { t.Fatalf("expected %s callback once, got %d", label, *calls) } @@ -99,6 +102,9 @@ func assertScopeLocalCallbackDeregisters( if err := runAfter(); err != nil { t.Fatalf("%s after deregister failed: %v", label, err) } + if err := FlushSubscribers(); err != nil { + t.Fatalf(scopeLocalFlushSubscribersFailed, err) + } if *calls != 1 { t.Fatalf("%s callback still fired after deregister: %d", label, *calls) } @@ -513,6 +519,9 @@ func TestPriorityMergeGlobalAndScopeLocal(t *testing.T) { if err != nil { t.Fatalf(scopeLocalToolCallExecuteFailed, err) } + if err := FlushSubscribers(); err != nil { + t.Fatalf(scopeLocalFlushSubscribersFailed, err) + } mu.Lock() defer mu.Unlock() if len(order) != 2 { @@ -564,6 +573,9 @@ func TestPriorityMergeGlobalBeforeScopeLocal(t *testing.T) { if err != nil { t.Fatalf(scopeLocalToolCallExecuteFailed, err) } + if err := FlushSubscribers(); err != nil { + t.Fatalf(scopeLocalFlushSubscribersFailed, err) + } mu.Lock() defer mu.Unlock() if len(order) != 2 { diff --git a/go/nemo_relay/tools_test.go b/go/nemo_relay/tools_test.go index e8c6b784b..14ee1769e 100644 --- a/go/nemo_relay/tools_test.go +++ b/go/nemo_relay/tools_test.go @@ -638,6 +638,9 @@ func TestToolMultipleGuardrailsPriorityOrder(t *testing.T) { if err != nil { t.Fatalf(toolCallExecuteFailed, err) } + if err := FlushSubscribers(); err != nil { + t.Fatalf(toolFlushSubscribersFailed, err) + } mu.Lock() defer mu.Unlock() diff --git a/python/tests/test_llm.py b/python/tests/test_llm.py index 4e2af4fc7..010b7501f 100644 --- a/python/tests/test_llm.py +++ b/python/tests/test_llm.py @@ -306,6 +306,7 @@ def sanitize_response(response, context): codec=codec, response_codec=codec, ) + await subscribers.flush_async() finally: guardrails.deregister_llm_sanitize_request("py_llm_builtin_context_request") guardrails.deregister_llm_sanitize_response("py_llm_builtin_context_response") diff --git a/python/tests/test_scope_local.py b/python/tests/test_scope_local.py index e0601dd35..b9539b600 100644 --- a/python/tests/test_scope_local.py +++ b/python/tests/test_scope_local.py @@ -226,6 +226,7 @@ def my_tool(args): with scope.scope("priority_scope", ScopeType.Agent) as handle: scope_local.register_tool_sanitize_request(handle, "sl_local_guard", 5, scope_local_sanitizer) await tools.execute("priority_tool", {"test": True}, my_tool) + await subscribers.flush_async() guardrails.deregister_tool_sanitize_request("sl_global_guard") @@ -253,6 +254,7 @@ def my_tool(args): with scope.scope("priority_scope2", ScopeType.Agent) as handle: scope_local.register_tool_sanitize_request(handle, "sl_local_guard2", 20, scope_local_sanitizer) await tools.execute("priority_tool2", {}, my_tool) + await subscribers.flush_async() guardrails.deregister_tool_sanitize_request("sl_global_guard2") From d52229f194baa9b1ecece7b62fcafd03207fe3f8 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Mon, 10 Aug 2026 13:45:54 -0400 Subject: [PATCH 4/5] fix: preserve queued lifecycle boundaries Signed-off-by: Will Killian --- crates/core/src/api/llm.rs | 47 ++++++++++----- crates/core/src/api/tool.rs | 6 +- crates/core/src/stream.rs | 58 +++++-------------- .../tests/integration/middleware_tests.rs | 34 ++++++++++- crates/node/tests/callback_error_tests.mjs | 2 + docs/about-nemo-relay/architecture.mdx | 2 +- python/tests/test_scope_local.py | 2 + 7 files changed, 88 insertions(+), 63 deletions(-) diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index 236246e24..62e91f17f 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -500,20 +500,30 @@ fn queue_llm_start_with_subscribers( ) -> Result<()> { ensure_runtime_owner()?; let scope_stack = handle.captured_scope_stack().clone(); - let (entries, agent_is_fresh, full_payloads_enabled) = { - let mut scope_guard = scope_stack.write().expect("scope stack lock poisoned"); - let scope_locals = scope_guard.collect_scope_local_registries(|registries| { - ®istries.llm_sanitize_request_guardrails - }); + let scope_locals = { + let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); + scope_guard + .collect_scope_local_registries(|registries| { + ®istries.llm_sanitize_request_guardrails + }) + .into_iter() + .cloned() + .collect::>() + }; + let (entries, full_payloads_enabled) = { + let scope_local_refs = scope_locals.iter().collect::>(); let context = global_context(); let state = context .read() .map_err(|error| FlowError::Internal(error.to_string()))?; - let entries = state.llm_sanitize_request_entries(&scope_locals); - let full_payloads_enabled = state.observability_full_payloads_enabled; - drop(state); - let agent_is_fresh = scope_guard.take_agent_freshness(handle.parent_uuid); - (entries, agent_is_fresh, full_payloads_enabled) + ( + state.llm_sanitize_request_entries(&scope_local_refs), + state.observability_full_payloads_enabled, + ) + }; + let agent_is_fresh = { + let mut scope_guard = scope_stack.write().expect("scope stack lock poisoned"); + scope_guard.take_agent_freshness(handle.parent_uuid) }; let observable_request = remove_observability_credential_headers(request.clone()); let queued_handle = handle.clone(); @@ -1153,6 +1163,7 @@ async fn llm_call_end_with_behavior( response_codec, timestamp, } = params; + let timestamp = timestamp.unwrap_or_else(Utc::now); ensure_runtime_owner()?; let (entries, subscribers) = { let scope_stack = handle.captured_scope_stack(); @@ -1186,7 +1197,7 @@ async fn llm_call_end_with_behavior( .handle(handle) .data(Json::Null) .metadata_opt(end_metadata) - .timestamp_opt(timestamp) + .timestamp(timestamp) .build(), ) }; @@ -1223,7 +1234,7 @@ async fn llm_call_end_with_behavior( .data_opt(payload.data) .metadata_opt(end_metadata) .annotated_response_opt(payload.annotated_response) - .timestamp_opt(timestamp) + .timestamp(timestamp) .build(), ) }) @@ -1281,6 +1292,7 @@ async fn emit_llm_end_without_output( lifecycle_subscribers: Option<&[EventSubscriberFn]>, ) -> Result<()> { ensure_runtime_owner()?; + let timestamp = Utc::now(); let (entries, subscribers) = { let scope_stack = handle.captured_scope_stack(); let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); @@ -1313,6 +1325,7 @@ async fn emit_llm_end_without_output( .handle(handle) .data(Json::Null) .metadata_opt(metadata.clone()) + .timestamp(timestamp) .build(), ) }; @@ -1355,7 +1368,15 @@ async fn emit_llm_end_without_output( global_context() .read() .map(|state| { - state.end_llm_handle(&queued_handle, data, metadata, annotated_response) + state.build_llm_end_event( + EndLlmHandleParams::builder() + .handle(&queued_handle) + .data_opt(data) + .metadata_opt(metadata) + .annotated_response_opt(annotated_response) + .timestamp(timestamp) + .build(), + ) }) .unwrap_or(event) }) diff --git a/crates/core/src/api/tool.rs b/crates/core/src/api/tool.rs index a84af3327..933003ebf 100644 --- a/crates/core/src/api/tool.rs +++ b/crates/core/src/api/tool.rs @@ -335,8 +335,8 @@ async fn tool_call_with_subscriber_snapshot( ) -> Result<(ToolHandle, Vec)> { 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| { ®istries.tool_sanitize_request_guardrails @@ -389,7 +389,6 @@ async fn tool_call_with_subscriber_snapshot( .collect::>(); (handle, event, marks) }; - let scope_stack = current_scope_stack(); let event_sanitizers = snapshot_event_sanitizers(&event, &scope_stack).unwrap_or_default(); let tool_name = handle.name.clone(); dispatch_transformed_event( @@ -518,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| { ®istries.tool_sanitize_response_guardrails @@ -570,7 +569,6 @@ async fn tool_call_end_with_pending_marks( )) }) .collect::>(); - let scope_stack = current_scope_stack(); let event_sanitizers = snapshot_event_sanitizers(&event, &scope_stack).unwrap_or_default(); let tool_name = params.handle.name.clone(); let result = params.result; diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index bed67944f..994fb4ab2 100644 --- a/crates/core/src/stream.rs +++ b/crates/core/src/stream.rs @@ -86,7 +86,6 @@ pub struct LlmStreamWrapper { chunk_index: u64, ended: bool, close_result: Option>, - finalization: Option>, terminal_result: Option>, } @@ -180,7 +179,6 @@ impl LlmStreamWrapper { chunk_index: 0, ended: false, close_result: None, - finalization: None, terminal_result: None, } } @@ -197,7 +195,7 @@ impl LlmStreamWrapper { &self.scope_stack } - fn finish(&mut self, background_thread: bool) { + fn finish(&mut self) { if self.ended { return; } @@ -213,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) { @@ -223,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) { @@ -233,20 +233,17 @@ 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, - termination: StreamTermination, - _background_thread: bool, - ) -> Option> { + /// Calls the finalizer and queues response sanitization and END publication. + fn emit_end_event(&mut self, metadata: Option, 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. @@ -370,7 +367,6 @@ impl LlmStreamWrapper { // or event sanitizers. The registered barrier keeps subscriber flushes // ordered behind this END event. let _ = subscriber_dispatcher::spawn_background_publication(finalize); - None } /// Emit a compact per-chunk receipt mark before collector processing. @@ -438,27 +434,6 @@ impl Stream for LlmStreamWrapper { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.as_mut().get_mut(); - // Retain support for an already-scheduled finalization task while the - // stream is being polled. - 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)), @@ -504,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() @@ -787,7 +757,7 @@ fn non_empty_object(object: Map) -> Option { impl Drop for LlmStreamWrapper { fn drop(&mut self) { - self.finish(true); + self.finish(); } } diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index 877cc6a0c..65427022b 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -4440,6 +4440,8 @@ async fn test_managed_llm_payload_sanitizers_are_queued_off_execution_path() { Box::pin(async move { request_started.notify_one(); request_release.notified().await; + let mut request = request; + request.content["sanitized"] = json!(true); Ok(Some(request)) }) } @@ -4460,13 +4462,23 @@ async fn test_managed_llm_payload_sanitizers_are_queued_off_execution_path() { Box::pin(async move { response_started.notify_one(); response_release.notified().await; + let mut response = response; + response["sanitized"] = json!(true); Ok(Some(response)) }) } }), ) .unwrap(); - register_subscriber("managed_queued_llm_observer", Arc::new(|_| {})).unwrap(); + let events = Arc::new(Mutex::new(Vec::::new())); + register_subscriber( + "managed_queued_llm_observer", + Arc::new({ + let events = Arc::clone(&events); + move |event| events.lock().unwrap().push(event.clone()) + }), + ) + .unwrap(); let call = tokio::spawn(async { llm_call_execute( @@ -4494,6 +4506,7 @@ async fn test_managed_llm_payload_sanitizers_are_queued_off_execution_path() { .expect("request sanitizer blocked managed provider execution") .expect("managed call task should join") .expect("managed call should succeed"); + let completed_at = Utc::now(); assert_eq!(result, json!({"response": "done"})); request_release.notify_one(); @@ -4504,6 +4517,25 @@ async fn test_managed_llm_payload_sanitizers_are_queued_off_execution_path() { .await .expect("managed response sanitizer did not start"); assert_flush_waits_for_pending_completion(|| response_release.notify_one()); + let events = events.lock().unwrap(); + let start = events + .iter() + .find(|event| { + event.name() == "managed-queued-llm" + && event.scope_category() == Some(ScopeCategory::Start) + }) + .expect("managed LLM START event should be published after flush"); + assert_eq!(start.input().unwrap()["content"]["sanitized"], true); + let end = events + .iter() + .find(|event| { + event.name() == "managed-queued-llm" + && event.scope_category() == Some(ScopeCategory::End) + }) + .expect("managed LLM END event should be published after flush"); + assert_eq!(end.data().unwrap()["sanitized"], true); + assert!(*end.timestamp() <= completed_at); + drop(events); deregister_llm_sanitize_request_guardrail("managed_queued_llm_request").unwrap(); deregister_llm_sanitize_response_guardrail("managed_queued_llm_response").unwrap(); deregister_subscriber("managed_queued_llm_observer").unwrap(); diff --git a/crates/node/tests/callback_error_tests.mjs b/crates/node/tests/callback_error_tests.mjs index 43a5b326b..2244fe8ac 100644 --- a/crates/node/tests/callback_error_tests.mjs +++ b/crates/node/tests/callback_error_tests.mjs @@ -17,6 +17,7 @@ const { __testClosedToolCallback, clearLastCallbackError, deregisterLlmSanitizeRequestGuardrail, + flushSubscribers, getLastCallbackError, llmCallExecute, registerLlmSanitizeRequestGuardrail, @@ -54,6 +55,7 @@ describe('callback error helpers', () => { const { traceparent, ...headers } = result.headers; assert.deepEqual({ ...result, headers }, { model: 'test-model', headers: {} }); assert.match(traceparent, /^00-[0-9a-f]{32}-[0-9a-f]{16}-01$/); + await flushSubscribers(); assert.match( getLastCallbackError() ?? '', /JS LLM sanitize request callback failed: failed to deserialize LlmRequest/i, diff --git a/docs/about-nemo-relay/architecture.mdx b/docs/about-nemo-relay/architecture.mdx index 03c28dd51..fb0a5e68f 100644 --- a/docs/about-nemo-relay/architecture.mdx +++ b/docs/about-nemo-relay/architecture.mdx @@ -118,7 +118,7 @@ model: Every managed tool or LLM call resolves the conditional and intercept middleware visible from the active scope before it executes. Payload sanitization and event publication are queued observability work: they do not delay the real callback -or result, and `flush_subscribers()` is the explicit delivery barrier. When the +or result, and `flush_subscribers` is the explicit delivery barrier. When the runtime emits an event, it records the active scope UUID as parentage. The scope stack changes as work opens and closes; the parent-linked event records remain available to subscribers. diff --git a/python/tests/test_scope_local.py b/python/tests/test_scope_local.py index b9539b600..4a8623bea 100644 --- a/python/tests/test_scope_local.py +++ b/python/tests/test_scope_local.py @@ -471,10 +471,12 @@ def my_tool(args): with scope.scope("persist_scope_1", ScopeType.Agent) as handle: scope_local.register_tool_sanitize_request(handle, "sl_persist_local", 2, lambda n, a: a) await tools.execute("persist_tool_1", {}, my_tool) + await subscribers.flush_async() # Global should still work after the scope-local scope ends execution_order.clear() await tools.execute("persist_tool_2", {}, my_tool) + await subscribers.flush_async() guardrails.deregister_tool_sanitize_request("sl_persist_global") From ca5acf55c490aeb680e1ef973f7e8eea40abd664 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Mon, 10 Aug 2026 14:04:33 -0400 Subject: [PATCH 5/5] fix: isolate provider trace context from observability Signed-off-by: Will Killian --- crates/core/src/api/llm.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index 62e91f17f..0388e1353 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -1638,10 +1638,11 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); snapshot_event_subscribers(scope_guard.collect_scope_local_subscribers())? }; + let observability_request = intercepted_request.clone(); inject_traceparent(&mut intercepted_request, handle.uuid)?; queue_llm_start_with_subscribers( &handle, - &intercepted_request, + &observability_request, annotated_request.clone(), request_codec.clone(), &lifecycle_subscribers, @@ -1848,10 +1849,11 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); snapshot_event_subscribers(scope_guard.collect_scope_local_subscribers())? }; + let observability_request = intercepted_request.clone(); inject_traceparent(&mut intercepted_request, handle.uuid)?; queue_llm_start_with_subscribers( &handle, - &intercepted_request, + &observability_request, annotated_request, request_codec.clone(), &lifecycle_subscribers,