feat(sanitizers)!: add mark and scope event sanitizer registries [1/6]

crates/core/src/api/llm.rs

@@ -21,7 +21,8 @@ use crate::api::scope::event;
 use crate::api::scope::{EmitMarkEventParams, ScopeHandle};
 use crate::api::shared::{
     ensure_runtime_owner, inject_dynamo_session_ids, metadata_with_otel_status,
-    resolve_parent_uuid, run_request_intercepts_with_codec, snapshot_event_subscribers,
+    resolve_parent_uuid, run_request_intercepts_with_codec, sanitize_event,
+    snapshot_event_subscribers,
 };
 use crate::codec::request::AnnotatedLlmRequest;
 use crate::codec::response::{AnnotatedLlmResponse, attach_estimated_cost_for_provider};
@@ -320,7 +321,9 @@ fn emit_llm_start_with_subscribers(
             .map_err(|error| FlowError::Internal(error.to_string()))?;
         state.build_llm_start_event(handle, Some(input), annotated_request)
     };
-    NemoRelayContextState::emit_event(&event, subscribers);
+    if let Some(event) = sanitize_event(event) {
+        NemoRelayContextState::emit_event(&event, subscribers);
+    }
     Ok(())
 }
 
@@ -346,7 +349,9 @@ fn emit_pending_request_marks(
             mark.category,
             mark.category_profile,
         ));
-        NemoRelayContextState::emit_event(&event, subscribers);
+        if let Some(event) = sanitize_event(event) {
+            NemoRelayContextState::emit_event(&event, subscribers);
+        }
     }
     Ok(())
 }
@@ -519,7 +524,9 @@ fn llm_call_end_with_behavior(
                 .build(),
         )
     };
-    NemoRelayContextState::emit_event(&event, &subscribers);
+    if let Some(event) = sanitize_event(event) {
+        NemoRelayContextState::emit_event(&event, &subscribers);
+    }
     if let Some(error) = decode_error
         && behavior.response_codec_errors_fatal
     {
@@ -550,7 +557,9 @@ fn emit_llm_end_without_output(
         let event = state.end_llm_handle(handle, handle.data.clone(), metadata, None);
         (event, subscribers)
     };
-    NemoRelayContextState::emit_event(&event, &subscribers);
+    if let Some(event) = sanitize_event(event) {
+        NemoRelayContextState::emit_event(&event, &subscribers);
+    }
     Ok(())
 }
 

crates/core/src/api/registry.rs

@@ -5,7 +5,7 @@
 //! intercepts, and subscribers.
 
 use crate::api::runtime::{
-    LlmConditionalFn, LlmExecutionFn, LlmRequestInterceptFn, LlmSanitizeRequestFn,
+    EventSanitizeFn, LlmConditionalFn, LlmExecutionFn, LlmRequestInterceptFn, LlmSanitizeRequestFn,
     LlmSanitizeResponseFn, LlmStreamExecutionFn, ToolConditionalFn, ToolExecutionFn,
     ToolInterceptFn, ToolSanitizeFn,
 };
@@ -467,6 +467,31 @@ macro_rules! scope_execution_registry_api {
     };
 }
 
+global_guardrail_registry_api!(
+    /// Register a global mark event sanitizer.
+    register_mark_sanitize_guardrail,
+    /// Deregister a global mark event sanitizer.
+    deregister_mark_sanitize_guardrail,
+    mark_sanitize_guardrails,
+    EventSanitizeFn
+);
+global_guardrail_registry_api!(
+    /// Register a global scope-start event sanitizer.
+    register_scope_sanitize_start_guardrail,
+    /// Deregister a global scope-start event sanitizer.
+    deregister_scope_sanitize_start_guardrail,
+    scope_sanitize_start_guardrails,
+    EventSanitizeFn
+);
+global_guardrail_registry_api!(
+    /// Register a global scope-end event sanitizer.
+    register_scope_sanitize_end_guardrail,
+    /// Deregister a global scope-end event sanitizer.
+    deregister_scope_sanitize_end_guardrail,
+    scope_sanitize_end_guardrails,
+    EventSanitizeFn
+);
+
 global_guardrail_registry_api!(
     /// Register a global tool sanitize-request guardrail.
     /// The guardrail rewrites only the tool input recorded on emitted start
@@ -579,6 +604,31 @@ global_execution_registry_api!(
     LlmStreamExecutionFn
 );
 
+scope_guardrail_registry_api!(
+    /// Register a scope-local mark event sanitizer.
+    scope_register_mark_sanitize_guardrail,
+    /// Deregister a scope-local mark event sanitizer.
+    scope_deregister_mark_sanitize_guardrail,
+    mark_sanitize_guardrails,
+    EventSanitizeFn
+);
+scope_guardrail_registry_api!(
+    /// Register a scope-local scope-start event sanitizer.
+    scope_register_scope_sanitize_start_guardrail,
+    /// Deregister a scope-local scope-start event sanitizer.
+    scope_deregister_scope_sanitize_start_guardrail,
+    scope_sanitize_start_guardrails,
+    EventSanitizeFn
+);
+scope_guardrail_registry_api!(
+    /// Register a scope-local scope-end event sanitizer.
+    scope_register_scope_sanitize_end_guardrail,
+    /// Deregister a scope-local scope-end event sanitizer.
+    scope_deregister_scope_sanitize_end_guardrail,
+    scope_sanitize_end_guardrails,
+    EventSanitizeFn
+);
+
 scope_guardrail_registry_api!(
     /// Register a scope-local tool sanitize-request guardrail.
     /// The guardrail rewrites only tool input emitted under the owning scope.

crates/core/src/api/runtime.rs

@@ -10,8 +10,8 @@ pub mod state;
 pub mod subscriber_dispatcher;
 
 pub use callbacks::{
-    EventSubscriberFn, LlmCollectorFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn,
-    LlmFinalizerFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestFn,
+    EventSanitizeFn, EventSubscriberFn, LlmCollectorFn, LlmConditionalFn, LlmExecutionFn,
+    LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestFn,
     LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, ToolConditionalFn,
     ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn,
 };

crates/core/src/api/runtime/callbacks.rs

@@ -14,13 +14,20 @@ use std::sync::Arc;
 
 use tokio_stream::Stream;
 
-use crate::api::event::Event;
+use crate::api::event::{Event, EventSanitizeFields};
 use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome};
 use crate::api::tool::ToolExecutionInterceptOutcome;
 use crate::codec::request::AnnotatedLlmRequest;
 use crate::error::Result;
 use crate::json::Json;
 
+/// Sanitize mutable observability fields on a fully constructed event.
+///
+/// The callback receives the current event as immutable context and the fields
+/// it may replace. Later callbacks observe fields returned by earlier entries.
+pub type EventSanitizeFn =
+    Arc<dyn Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync>;
+
 /// Sanitize a tool request payload before the runtime records it.
 ///
 /// Tool sanitize callbacks are used only for observability payloads. They can

crates/core/src/api/runtime/state.rs

@@ -21,13 +21,14 @@ use crate::api::llm::{CreateLlmHandleParams, EndLlmHandleParams};
 use crate::api::llm::{LlmHandle, LlmRequest};
 use crate::api::registry::{ExecutionIntercept, Guardrail, Intercept};
 use crate::api::runtime::callbacks::{
-    EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, LlmRequestInterceptFn,
-    LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn,
-    LlmStreamExecutionRegistryRefs, ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn,
-    ToolExecutionOutcomeNextFn, ToolInterceptFn, ToolSanitizeFn,
+    EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn,
+    LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn,
+    LlmStreamExecutionNextFn, LlmStreamExecutionRegistryRefs, ToolConditionalFn, ToolExecutionFn,
+    ToolExecutionNextFn, ToolExecutionOutcomeNextFn, ToolInterceptFn, ToolSanitizeFn,
 };
 use crate::api::runtime::subscriber_dispatcher;
 use crate::api::scope::{CreateScopeHandleParams, EndScopeHandleParams, ScopeHandle, ScopeType};
+use crate::api::shared::sanitize_event;
 use crate::api::tool::ToolHandle;
 use crate::api::tool::{
     CreateToolHandleParams, EndToolHandleParams, ToolExecutionInterceptOutcome,
@@ -49,6 +50,12 @@ use uuid::Uuid;
 /// process. It contains global middleware registries, lifecycle subscribers,
 /// and arbitrary extension slots used by bindings or integrations.
 pub struct NemoRelayContextState {
+    /// Global mark event field sanitizers.
+    pub(crate) mark_sanitize_guardrails: SortedRegistry<Guardrail<EventSanitizeFn>>,
+    /// Global scope-start event field sanitizers.
+    pub(crate) scope_sanitize_start_guardrails: SortedRegistry<Guardrail<EventSanitizeFn>>,
+    /// Global scope-end event field sanitizers.
+    pub(crate) scope_sanitize_end_guardrails: SortedRegistry<Guardrail<EventSanitizeFn>>,
     /// Global tool request sanitizers applied to emitted tool-start payloads.
     pub(crate) tool_sanitize_request_guardrails: SortedRegistry<Guardrail<ToolSanitizeFn>>,
     /// Global tool response sanitizers applied to emitted tool-end payloads.
@@ -86,6 +93,9 @@ impl NemoRelayContextState {
     /// extensions.
     pub fn new() -> Self {
         Self {
+            mark_sanitize_guardrails: SortedRegistry::new(),
+            scope_sanitize_start_guardrails: SortedRegistry::new(),
+            scope_sanitize_end_guardrails: SortedRegistry::new(),
             tool_sanitize_request_guardrails: SortedRegistry::new(),
             tool_sanitize_response_guardrails: SortedRegistry::new(),
             tool_conditional_execution_guardrails: SortedRegistry::new(),
@@ -571,7 +581,9 @@ impl NemoRelayContextState {
             EventCategory::from(handle.scope_type),
             None,
         ));
-        Self::emit_event(&event, subscribers);
+        if let Some(event) = sanitize_event(event) {
+            Self::emit_event(&event, subscribers);
+        }
         handle
     }
 
@@ -594,7 +606,32 @@ impl NemoRelayContextState {
             EventCategory::from(handle.scope_type),
             None,
         ));
-        Self::emit_event(&event, subscribers);
+        if let Some(event) = sanitize_event(event) {
+            Self::emit_event(&event, subscribers);
+        }
+    }
+
+    /// Snapshot event sanitizer entries in priority order.
+    pub(crate) fn event_sanitize_entries(
+        global: &SortedRegistry<Guardrail<EventSanitizeFn>>,
+        scope_locals: &[&SortedRegistry<Guardrail<EventSanitizeFn>>],
+    ) -> Vec<Guardrail<EventSanitizeFn>> {
+        merge_guardrail_entries(global, scope_locals)
+            .into_iter()
+            .cloned()
+            .collect()
+    }
+
+    /// Apply an event sanitizer snapshot to the mutable observability fields.
+    pub(crate) fn event_sanitize_snapshot_chain(
+        mut event: Event,
+        entries: &[Guardrail<EventSanitizeFn>],
+    ) -> Event {
+        for entry in entries {
+            let fields = (entry.payload)(&event, event.sanitize_fields());
+            event.apply_sanitize_fields(fields);
+        }
+        event
     }
 
     /// Snapshot tool request sanitizers in priority order.

crates/core/src/api/scope.rs

@@ -7,7 +7,9 @@ use crate::api::runtime::global_context;
 use crate::api::runtime::{
     current_scope_stack, task_scope_push, task_scope_remove, task_scope_top,
 };
-use crate::api::shared::{ensure_runtime_owner, resolve_parent_uuid, snapshot_event_subscribers};
+use crate::api::shared::{
+    ensure_runtime_owner, resolve_parent_uuid, sanitize_event, snapshot_event_subscribers,
+};
 use crate::error::{FlowError, Result};
 use crate::json::Json;
 use chrono::{DateTime, Utc};
@@ -229,8 +231,11 @@ pub fn push_scope(params: PushScopeParams<'_>) -> Result<ScopeHandle> {
         let event = state.build_scope_start_event(&handle, params.input);
         (handle, event, subscribers)
     };
+    let event = sanitize_event(event);
     task_scope_push(handle.clone());
-    NemoRelayContextState::emit_event(&event, &subscribers);
+    if let Some(event) = event {
+        NemoRelayContextState::emit_event(&event, &subscribers);
+    }
     Ok(handle)
 }
 
@@ -287,9 +292,12 @@ pub fn pop_scope(params: PopScopeParams<'_>) -> Result<()> {
         );
         (scope, event, subscribers)
     };
+    let event = sanitize_event(event);
     let removed = task_scope_remove(params.handle_uuid)?;
     debug_assert_eq!(removed.uuid, scope.uuid);
-    NemoRelayContextState::emit_event(&event, &subscribers);
+    if let Some(event) = event {
+        NemoRelayContextState::emit_event(&event, &subscribers);
+    }
     Ok(())
 }
 
@@ -342,6 +350,8 @@ pub fn event(params: EmitMarkEventParams<'_>) -> Result<()> {
         ));
         (event, subscribers)
     };
-    NemoRelayContextState::emit_event(&event, &subscribers);
+    if let Some(event) = sanitize_event(event) {
+        NemoRelayContextState::emit_event(&event, &subscribers);
+    }
     Ok(())
 }

crates/core/src/api/shared.rs

@@ -5,9 +5,10 @@ use std::sync::Arc;
 
 use uuid::Uuid;
 
+use crate::api::event::{Event, ScopeCategory};
 use crate::api::llm::LlmRequest;
-use crate::api::runtime::EventSubscriberFn;
 use crate::api::runtime::global_context;
+use crate::api::runtime::{EventSubscriberFn, NemoRelayContextState, ScopeStackHandle};
 use crate::api::runtime::{current_scope_stack, task_scope_top};
 use crate::api::scope::ScopeHandle;
 use crate::api::scope::ScopeType;
@@ -40,6 +41,58 @@ pub(crate) fn snapshot_event_subscribers(
     Ok(state.collect_event_subscribers(&scope_local_subscribers))
 }
 
+/// Apply the event sanitizer chain visible on the current scope stack.
+pub(crate) fn sanitize_event(event: Event) -> Option<Event> {
+    sanitize_event_with_scope_stack(event, &current_scope_stack())
+}
+
+/// Apply the event sanitizer chain visible on a captured scope stack.
+pub(crate) fn sanitize_event_with_scope_stack(
+    event: Event,
+    scope_stack: &ScopeStackHandle,
+) -> Option<Event> {
+    let entries = {
+        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
+        let context = global_context();
+        let state = match context.read() {
+            Ok(state) => state,
+            Err(_) => return None,
+        };
+        match &event {
+            Event::Mark(_) => {
+                let locals = scope_guard.collect_scope_local_registries(|registries| {
+                    &registries.mark_sanitize_guardrails
+                });
+                NemoRelayContextState::event_sanitize_entries(
+                    &state.mark_sanitize_guardrails,
+                    &locals,
+                )
+            }
+            Event::Scope(scope) if scope.scope_category == ScopeCategory::Start => {
+                let locals = scope_guard.collect_scope_local_registries(|registries| {
+                    &registries.scope_sanitize_start_guardrails
+                });
+                NemoRelayContextState::event_sanitize_entries(
+                    &state.scope_sanitize_start_guardrails,
+                    &locals,
+                )
+            }
+            Event::Scope(_) => {
+                let locals = scope_guard.collect_scope_local_registries(|registries| {
+                    &registries.scope_sanitize_end_guardrails
+                });
+                NemoRelayContextState::event_sanitize_entries(
+                    &state.scope_sanitize_end_guardrails,
+                    &locals,
+                )
+            }
+        }
+    };
+    Some(NemoRelayContextState::event_sanitize_snapshot_chain(
+        event, &entries,
+    ))
+}
+
 pub(crate) fn ensure_runtime_owner() -> Result<()> {
     ensure_process_runtime_owner()
 }