From 8f6daccf4fd8cf5b70caae5c026dd9f75c9dead6 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Tue, 4 Aug 2026 13:25:34 -0400 Subject: [PATCH 1/2] fix: queue managed sanitizer publication Signed-off-by: Will Killian --- .agents/skills/add-middleware/SKILL.md | 10 +- AGENTS.md | 2 +- crates/core/src/api/llm.rs | 274 +++++++++++----- crates/core/src/api/tool.rs | 83 +++-- crates/core/src/stream.rs | 69 ++-- .../tests/integration/middleware_tests.rs | 302 ++++++++++++++++-- crates/core/tests/unit/llm_api_tests.rs | 2 +- crates/node/tests/llm_tests.mjs | 2 + docs/about-nemo-relay/architecture.mdx | 2 +- docs/about-nemo-relay/concepts/middleware.mdx | 80 ++--- 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 + 15 files changed, 641 insertions(+), 206 deletions(-) diff --git a/.agents/skills/add-middleware/SKILL.md b/.agents/skills/add-middleware/SKILL.md index 40a38bbfd..a362e36a7 100644 --- a/.agents/skills/add-middleware/SKILL.md +++ b/.agents/skills/add-middleware/SKILL.md @@ -41,11 +41,13 @@ Decide these before editing code: Refer to `docs/about-nemo-relay/concepts/middleware.mdx` for the full diagrams. - **Tool execute**: - conditional guardrails -> request intercepts -> sanitize request (for events) - | execution intercept chain(callable) -> sanitize response + conditional guardrails -> request intercepts -> queue sanitize request (for events) + -> execution intercept chain(callable) -> queue sanitize response - **LLM execute**: - conditional guardrails -> request intercepts -> sanitize request (for events) - | execution intercept chain(callable) -> sanitize response + conditional guardrails -> request intercepts -> queue sanitize request (for events) + -> execution intercept chain(callable) -> queue sanitize response +- **Queued publication**: + sanitize payload -> sanitize event fields -> subscriber/exporter delivery - **Mark and scope events**: specialized tool or LLM sanitizer (when applicable) -> mark or scope event sanitizer -> subscriber and exporter dispatch diff --git a/AGENTS.md b/AGENTS.md index d642699f5..6acead826 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -177,7 +177,7 @@ preserve. - Middleware is priority-ordered after merging global and visible scope-local entries. - Intercepts change the real execution path. Request intercepts rewrite the request. Execution intercepts wrap or replace the callback. Stream execution intercepts handle streaming lifecycle behavior. - Guardrails either block execution or sanitize emitted observability payloads. Sanitize guardrails do not rewrite the real callback arguments or return value. -- Managed execution order is conditional guardrails, request intercepts, sanitize-request guardrails for start events, execution intercepts, callback execution, then sanitize-response guardrails for end events. +- Managed execution runs conditional guardrails and request intercepts, queues sanitize-request guardrails for the start event, runs execution intercepts and callback execution, then queues sanitize-response guardrails for the end event. Payload sanitizers run on the observability dispatcher and do not delay the application callback or result. - Events use ATOF `0.1` as the canonical event format. Scope events use start/end pairs; mark events record runtime checkpoints. - LLM and tool event metadata belongs in the category profile, such as `model_name`, `tool_call_id`, and custom `subtype` fields. - Exporters can transform runtime events to ATIF trajectories, OpenTelemetry traces, or OpenInference-compatible output. Root scope identity is used to isolate concurrent agents. diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index f398ebe93..6a98bf7ec 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -415,6 +415,7 @@ fn limit_annotated_request_history_to_current_user_turn( ) } +#[cfg(test)] async fn emit_llm_start_with_subscribers( handle: &LlmHandle, request: &LlmRequest, @@ -479,6 +480,83 @@ 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) = { + 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); + drop(state); + let agent_is_fresh = scope_guard.take_agent_freshness(handle.parent_uuid); + (entries, agent_is_fresh) + }; + 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 !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 @@ -830,7 +908,6 @@ pub fn llm_call(params: LlmCallParams<'_>) -> Result { #[derive(Clone, Copy)] struct LlmCallEndBehavior { - response_codec_errors_fatal: bool, attach_estimated_cost: bool, } @@ -996,7 +1073,6 @@ pub fn llm_call_end(params: LlmCallEndParams<'_>) -> Result<()> { response_codec, &entries, LlmCallEndBehavior { - response_codec_errors_fatal: false, attach_estimated_cost: false, }, ) @@ -1065,41 +1141,67 @@ 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(); + 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 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(), + ) + }) + }), + event_sanitizers, + &subscribers, + scope_stack, + ); + Ok(()) } #[cfg(test)] @@ -1167,46 +1269,66 @@ 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.end_llm_handle(handle, None, metadata.clone(), None) }; - queue_sanitized_event_with_scope_stack(event, &subscribers, handle.captured_scope_stack()); + let scope_stack = handle.captured_scope_stack().clone(); + 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 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) + }) + }), + event_sanitizers, + &subscribers, + scope_stack, + ); Ok(()) } @@ -1330,10 +1452,11 @@ impl Drop for ManagedLlmCompletion { /// Execute an LLM call through the managed middleware pipeline. /// -/// This runs conditional-execution guardrails, request intercepts, and -/// sanitize-request guardrails, emits the LLM-start event, then runs execution -/// intercepts, the provider callback when it is not replaced, and -/// sanitize-response guardrails in the runtime-defined order. +/// This runs conditional-execution guardrails and request intercepts, queues +/// sanitize-request guardrails and the LLM-start event, then runs execution +/// intercepts and the provider callback when it is not replaced. +/// Sanitize-response guardrails and the LLM-end event are queued before the +/// real response is returned. /// /// # Parameters /// - `name`: Logical provider or model family name recorded on emitted events. @@ -1360,11 +1483,11 @@ impl Drop for ManagedLlmCompletion { /// execution intercepts, codecs, or the callback itself. /// /// # Notes -/// The LLM-start event is emitted before execution intercepts run. Before -/// sanitize-request guardrails run, the runtime removes standard credential -/// headers from the event-only request copy; the request passed to execution is -/// unchanged. When execution fails after that point, the runtime still emits an -/// LLM-end event without an output payload. +/// The LLM-start event is queued before execution intercepts run. Before the +/// queued sanitize-request guardrails run, the runtime removes standard +/// credential headers from the event-only request copy; the request passed to +/// execution is unchanged. When execution fails after that point, the runtime +/// still queues an LLM-end event without an output payload. /// /// Response codecs enrich observability output only and do not change the /// value returned to the caller. @@ -1459,14 +1582,13 @@ 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( + queue_llm_start_with_subscribers( &handle, &intercepted_request, annotated_request.clone(), request_codec.clone(), &lifecycle_subscribers, - ) - .await?; + )?; emit_pending_request_marks(&handle, pending_marks, &lifecycle_subscribers).await?; handle .optimization_recorder @@ -1512,7 +1634,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), @@ -1568,10 +1689,10 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { /// execution intercepts, stream callbacks, codecs, or the provider callback. /// /// # Notes -/// The LLM-start event is emitted before stream execution intercepts run. -/// Before sanitize-request guardrails run, the runtime removes standard -/// credential headers from the event-only request copy; the request passed to -/// stream execution is unchanged. +/// The LLM-start event is queued before stream execution intercepts run. +/// Before the queued sanitize-request guardrails run, the runtime removes +/// standard credential headers from the event-only request copy; the request +/// passed to stream execution is unchanged. /// /// The returned stream emits chunk-level results while the runtime defers the /// LLM-end event until the collector and finalizer complete. @@ -1668,14 +1789,13 @@ 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( + queue_llm_start_with_subscribers( &handle, &intercepted_request, annotated_request, request_codec.clone(), &lifecycle_subscribers, - ) - .await?; + )?; 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..825e62f3b 100644 --- a/crates/core/src/api/tool.rs +++ b/crates/core/src/api/tool.rs @@ -356,12 +356,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 +372,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,7 +394,27 @@ 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, + ); for mark in marks { queue_sanitized_event(mark, &subscribers); } @@ -526,19 +541,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 +549,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,7 +574,35 @@ 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, + ); for mark in marks { queue_sanitized_event(mark, subscribers); } @@ -652,9 +682,10 @@ impl Drop for ManagedToolCompletion { /// Execute a tool call through the managed middleware pipeline. /// -/// This runs conditional-execution guardrails, request intercepts, -/// sanitize-request guardrails, execution intercepts, the tool callback, and -/// sanitize-response guardrails in the runtime-defined order. +/// This runs conditional-execution guardrails and request intercepts, queues +/// sanitize-request guardrails for observability, then runs execution +/// intercepts and the tool callback. Sanitize-response guardrails are queued +/// before the real result is returned. /// /// # Parameters /// - `name`: Tool name recorded on emitted lifecycle events. diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index 4679b2751..9b4dddea4 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; @@ -59,9 +61,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 @@ -212,6 +214,9 @@ impl LlmStreamWrapper { self.inner.terminalize(); let metadata = metadata_with_otel_status(self.metadata.clone(), status_code, status_message); + self.handle + .optimization_recorder + .close_for_finalization(None); self.finalization = self.emit_end_event(metadata, interrupted, false); } @@ -222,6 +227,9 @@ impl LlmStreamWrapper { self.ended = true; self.inner.terminalize(); let metadata = metadata_with_otel_error(self.metadata.clone(), error); + self.handle + .optimization_recorder + .close_for_finalization(None); self.finalization = self.emit_end_event(metadata, interrupted, false); } @@ -233,12 +241,13 @@ impl LlmStreamWrapper { &mut self, metadata: Option, interrupted: bool, - 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, @@ -314,9 +323,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, } }; @@ -337,22 +354,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. @@ -420,10 +427,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 8bc84a9cc..6c96d6a5b 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}; @@ -3439,7 +3440,6 @@ 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(); ready(args) }), ) @@ -3451,7 +3451,6 @@ 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(); ready(args) }), ) @@ -3485,7 +3484,6 @@ 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(); ready(result) }), ) @@ -3497,7 +3495,6 @@ 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(); ready(result) }), ) @@ -3519,6 +3516,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, &[ @@ -3614,7 +3612,6 @@ 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(); ready(Some(request)) }), ) @@ -3626,7 +3623,6 @@ 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(); ready(Some(request)) }), ) @@ -3683,7 +3679,6 @@ 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(); ready(Some(response)) }), ) @@ -3695,7 +3690,6 @@ 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(); ready(Some(response)) }), ) @@ -3760,6 +3754,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, &[ @@ -3902,25 +3897,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); @@ -4380,6 +4369,267 @@ 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(); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + register_subscriber( + "managed_queued_llm_observer", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .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()); + + let lifecycle = events + .lock() + .unwrap() + .iter() + .filter(|event| event.name() == "managed-queued-llm") + .map(|event| event.scope_category()) + .collect::>(); + assert_eq!( + lifecycle, + [Some(ScopeCategory::Start), Some(ScopeCategory::End)] + ); + + 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.category().map(|category| category.as_str()) == Some("llm") + && event.scope_category() == Some(ScopeCategory::End) + }) + .cloned() + .expect("stream END event should be published after flush"); + assert!( + *end.timestamp() <= terminal_timestamp, + "stream END timestamp must be captured before queued sanitization" + ); + + 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/crates/core/tests/unit/llm_api_tests.rs b/crates/core/tests/unit/llm_api_tests.rs index ee989abc9..4e52adb6d 100644 --- a/crates/core/tests/unit/llm_api_tests.rs +++ b/crates/core/tests/unit/llm_api_tests.rs @@ -1179,8 +1179,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 a2254af50..a5d00c97c 100644 --- a/crates/node/tests/llm_tests.mjs +++ b/crates/node/tests/llm_tests.mjs @@ -501,6 +501,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); @@ -611,6 +612,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/docs/about-nemo-relay/architecture.mdx b/docs/about-nemo-relay/architecture.mdx index 8af4a69d2..7766f3cbc 100644 --- a/docs/about-nemo-relay/architecture.mdx +++ b/docs/about-nemo-relay/architecture.mdx @@ -122,7 +122,7 @@ The plugin system installs reusable runtime components from configuration. A plu ### Event Emission -The runtime emits structured events for scopes, tools, LLMs, and named marks. Those events are the canonical record of runtime behavior. Native Rust, Python, Node.js, and FFI event-producing APIs enqueue subscriber work and return without waiting for subscriber callbacks or exporter work. +The runtime emits structured events for scopes, tools, LLMs, and named marks. Those events are the canonical record of runtime behavior. Native Rust, Python, Node.js, and FFI event-producing APIs enqueue payload sanitization, event sanitization, and subscriber work, then return without waiting for that observability work. ### Subscribers and Exporters diff --git a/docs/about-nemo-relay/concepts/middleware.mdx b/docs/about-nemo-relay/concepts/middleware.mdx index 19d583755..77cba157d 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, @@ -155,9 +157,11 @@ object. Guardrails are middleware that block execution or sanitize observability payloads. -Sanitizers are pure transformations. Do not use a sanitizer callback for -stateful side effects, such as metrics, logging, mutation, or I/O; Relay only -guarantees the transformed observability payload. +Sanitizers are transformations. Do not use a sanitizer callback for stateful +side effects such as metrics, logging, or application mutation; Relay only +guarantees the transformed observability payload. A sanitizer may await an +external transformation service, but slow calls delay queued observability and +subscriber flushes even though they do not delay managed execution. ### Conditional Execution @@ -199,11 +203,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 @@ -237,38 +242,37 @@ sequenceDiagram Cond-->>Caller: reject execution else allowed Runtime->>Req: rewrite the real request - Runtime->>San: sanitize emitted start payload - Runtime->>EventSan: sanitize start event fields - Runtime->>Dispatch: enqueue start event before execution - Dispatch-->>Consumers: deliver start event later + Runtime->>Dispatch: enqueue start-event copy before execution Runtime->>Exec: wrap execution Exec->>Callback: invoke callback Callback-->>Exec: return real result Exec-->>Runtime: continue - Runtime->>San: sanitize emitted end payload - Runtime->>EventSan: sanitize end event fields - Runtime->>Dispatch: enqueue end event - Dispatch-->>Consumers: deliver end event later + Runtime->>Dispatch: enqueue end-event copy Runtime-->>Caller: return real result + Dispatch->>San: sanitize start and end payloads in FIFO order + Dispatch->>EventSan: sanitize event fields + Dispatch-->>Consumers: deliver events later end ``` 1. Conditional-execution guardrails 2. Request intercepts -3. Tool or LLM sanitize-request guardrails -4. Scope-start event sanitizers and start-event emission -5. Execution intercepts -6. The real callback, unless an execution intercept replaces it -7. Tool or LLM sanitize-response guardrails -8. Scope-end event sanitizers and end-event emission - -For streaming LLM flows, the same pre-execution order applies: the runtime -applies `sanitize-request` guardrails and emits the LLM start event before the +3. Queue the start-event copy and its sanitize-request chain +4. Execution intercepts +5. The real callback, unless an execution intercept replaces it +6. Queue the end-event copy and its sanitize-response chain +7. Return the real result without waiting for either sanitizer chain + +On the publication path, the dispatcher preserves start/end ordering. It runs +the tool or LLM payload sanitizer first, then the matching scope-event +sanitizer, and finally delivers the event to subscribers and exporters. + +For streaming LLM flows, the runtime queues the LLM start-event copy before the stream execution intercept chain runs. Stream execution intercepts are the -execution family for streaming provider callbacks. The runtime then collects -chunks and finalizes the stream before `sanitize-response` guardrails rewrite -the emitted end-event payload and scope-end event sanitizers run at items 7 and -8. +execution family for streaming provider callbacks. The runtime collects chunks +and invokes the finalizer, then queues response and event sanitization without +delaying observable stream termination. A subscriber flush waits for that +queued end-event work. This ordering is what makes the semantic split between intercepts and guardrails important: 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 0065ed786..2619b8291 100644 --- a/python/tests/test_llm.py +++ b/python/tests/test_llm.py @@ -302,6 +302,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 1d692314791eb1dcd7117a92f12ae58eccd22403 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Tue, 4 Aug 2026 15:36:55 -0400 Subject: [PATCH 2/2] test: await queued llm sanitizer callbacks Signed-off-by: Will Killian --- crates/node/tests/llm_tests.mjs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/node/tests/llm_tests.mjs b/crates/node/tests/llm_tests.mjs index a5d00c97c..1cba57ff0 100644 --- a/crates/node/tests/llm_tests.mjs +++ b/crates/node/tests/llm_tests.mjs @@ -396,9 +396,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', );