diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index ab1a99e9f..a3d823fc4 100644 --- a/crates/core/src/api/llm.rs +++ b/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(()) } diff --git a/crates/core/src/api/registry.rs b/crates/core/src/api/registry.rs index 998096eec..72199095c 100644 --- a/crates/core/src/api/registry.rs +++ b/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. diff --git a/crates/core/src/api/runtime.rs b/crates/core/src/api/runtime.rs index 2351ae352..7d5a12882 100644 --- a/crates/core/src/api/runtime.rs +++ b/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, }; diff --git a/crates/core/src/api/runtime/callbacks.rs b/crates/core/src/api/runtime/callbacks.rs index d3790fbb7..ad81932d0 100644 --- a/crates/core/src/api/runtime/callbacks.rs +++ b/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 EventSanitizeFields + Send + Sync>; + /// Sanitize a tool request payload before the runtime records it. /// /// Tool sanitize callbacks are used only for observability payloads. They can diff --git a/crates/core/src/api/runtime/state.rs b/crates/core/src/api/runtime/state.rs index e5c134756..43f63dfd7 100644 --- a/crates/core/src/api/runtime/state.rs +++ b/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>, + /// Global scope-start event field sanitizers. + pub(crate) scope_sanitize_start_guardrails: SortedRegistry>, + /// Global scope-end event field sanitizers. + pub(crate) scope_sanitize_end_guardrails: SortedRegistry>, /// Global tool request sanitizers applied to emitted tool-start payloads. pub(crate) tool_sanitize_request_guardrails: SortedRegistry>, /// 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>, + scope_locals: &[&SortedRegistry>], + ) -> Vec> { + 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], + ) -> 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. diff --git a/crates/core/src/api/scope.rs b/crates/core/src/api/scope.rs index 3a93c8fa1..b00a97fbf 100644 --- a/crates/core/src/api/scope.rs +++ b/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 { 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(()) } diff --git a/crates/core/src/api/shared.rs b/crates/core/src/api/shared.rs index 28767174c..d9145d434 100644 --- a/crates/core/src/api/shared.rs +++ b/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 { + sanitize_event_with_scope_stack(event, ¤t_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 { + 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| { + ®istries.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| { + ®istries.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| { + ®istries.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() } diff --git a/crates/core/src/api/tool.rs b/crates/core/src/api/tool.rs index 2755f32df..3d60c7f80 100644 --- a/crates/core/src/api/tool.rs +++ b/crates/core/src/api/tool.rs @@ -11,7 +11,7 @@ use crate::api::runtime::{EventSubscriberFn, ToolExecutionNextFn}; use crate::api::scope::event; use crate::api::scope::{EmitMarkEventParams, ScopeHandle}; use crate::api::shared::{ - ensure_runtime_owner, metadata_with_otel_status, resolve_parent_uuid, + ensure_runtime_owner, metadata_with_otel_status, resolve_parent_uuid, sanitize_event, snapshot_event_subscribers, }; use crate::error::{FlowError, Result}; @@ -252,7 +252,9 @@ fn tool_call_with_subscriber_snapshot( let event = state.build_tool_start_event(&handle, Some(sanitized_args)); (handle, event) }; - NemoRelayContextState::emit_event(&event, &subscribers); + if let Some(event) = sanitize_event(event) { + NemoRelayContextState::emit_event(&event, &subscribers); + } Ok((handle, subscribers)) } @@ -353,8 +355,11 @@ fn tool_call_end_with_pending_marks( mark.category_profile, )) }) + .filter_map(sanitize_event) .collect::>(); - NemoRelayContextState::emit_event(&event, subscribers); + if let Some(event) = sanitize_event(event) { + NemoRelayContextState::emit_event(&event, subscribers); + } for mark in marks { NemoRelayContextState::emit_event(&mark, subscribers); } @@ -374,7 +379,9 @@ fn emit_tool_end_without_output( .map_err(|error| FlowError::Internal(error.to_string()))?; state.end_tool_handle(handle, handle.data.clone(), metadata) }; - NemoRelayContextState::emit_event(&event, lifecycle_subscribers); + if let Some(event) = sanitize_event(event) { + NemoRelayContextState::emit_event(&event, lifecycle_subscribers); + } Ok(()) } diff --git a/crates/core/src/context/registries.rs b/crates/core/src/context/registries.rs index 2a0d2fde9..bf59827a9 100644 --- a/crates/core/src/context/registries.rs +++ b/crates/core/src/context/registries.rs @@ -11,7 +11,7 @@ use std::collections::HashMap; use crate::api::registry::{ExecutionIntercept, Guardrail, Intercept}; use crate::api::runtime::{ - EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmRequestInterceptFn, + EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn, ToolConditionalFn, ToolExecutionFn, ToolInterceptFn, ToolSanitizeFn, }; @@ -24,6 +24,12 @@ use crate::registry::SortedRegistry; /// registries when the runtime resolves the effective middleware chain for a /// tool or LLM call executed inside that scope. pub(crate) struct ScopeLocalRegistries { + /// Mark event field sanitizers. + pub(crate) mark_sanitize_guardrails: SortedRegistry>, + /// Scope-start event field sanitizers. + pub(crate) scope_sanitize_start_guardrails: SortedRegistry>, + /// Scope-end event field sanitizers. + pub(crate) scope_sanitize_end_guardrails: SortedRegistry>, /// Tool request sanitizers applied to emitted tool-start payloads. pub(crate) tool_sanitize_request_guardrails: SortedRegistry>, /// Tool response sanitizers applied to emitted tool-end payloads. @@ -59,6 +65,9 @@ impl ScopeLocalRegistries { /// intercepts, or subscribers. pub(crate) 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(), diff --git a/crates/core/src/plugin.rs b/crates/core/src/plugin.rs index 196ddfb20..39a124ac3 100644 --- a/crates/core/src/plugin.rs +++ b/crates/core/src/plugin.rs @@ -23,17 +23,20 @@ use crate::api::registry::{ deregister_llm_conditional_execution_guardrail, deregister_llm_execution_intercept, deregister_llm_request_intercept, deregister_llm_sanitize_request_guardrail, deregister_llm_sanitize_response_guardrail, deregister_llm_stream_execution_intercept, - deregister_tool_conditional_execution_guardrail, deregister_tool_execution_intercept, - deregister_tool_request_intercept, deregister_tool_sanitize_request_guardrail, - deregister_tool_sanitize_response_guardrail, register_llm_conditional_execution_guardrail, - register_llm_execution_intercept, register_llm_request_intercept, - register_llm_sanitize_request_guardrail, register_llm_sanitize_response_guardrail, - register_llm_stream_execution_intercept, register_tool_conditional_execution_guardrail, + deregister_mark_sanitize_guardrail, deregister_scope_sanitize_end_guardrail, + deregister_scope_sanitize_start_guardrail, deregister_tool_conditional_execution_guardrail, + deregister_tool_execution_intercept, deregister_tool_request_intercept, + deregister_tool_sanitize_request_guardrail, deregister_tool_sanitize_response_guardrail, + register_llm_conditional_execution_guardrail, register_llm_execution_intercept, + register_llm_request_intercept, register_llm_sanitize_request_guardrail, + register_llm_sanitize_response_guardrail, register_llm_stream_execution_intercept, + register_mark_sanitize_guardrail, register_scope_sanitize_end_guardrail, + register_scope_sanitize_start_guardrail, register_tool_conditional_execution_guardrail, register_tool_execution_intercept, register_tool_request_intercept, register_tool_sanitize_request_guardrail, register_tool_sanitize_response_guardrail, }; use crate::api::runtime::{ - EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmRequestInterceptFn, + EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn, ToolConditionalFn, ToolExecutionFn, ToolInterceptFn, ToolSanitizeFn, }; @@ -320,6 +323,89 @@ impl PluginRegistrationContext { Ok(()) } + /// Registers a mark event sanitizer and records its rollback closure. + pub fn register_mark_sanitize_guardrail( + &mut self, + name: &str, + priority: i32, + callback: EventSanitizeFn, + ) -> Result<()> { + let qualified_name = self.qualify_name(name); + register_mark_sanitize_guardrail(&qualified_name, priority, callback) + .map_err(|err| PluginError::RegistrationFailed(format!("mark sanitizer: {err}")))?; + let name_owned = qualified_name; + self.registrations.push(PluginRegistration::new( + "plugin", + name_owned.clone(), + Box::new(move || { + deregister_mark_sanitize_guardrail(&name_owned) + .map(|_| ()) + .map_err(|err| { + PluginError::RegistrationFailed(format!( + "mark sanitizer deregistration failed: {err}" + )) + }) + }), + )); + Ok(()) + } + + /// Registers a scope-start event sanitizer and records its rollback closure. + pub fn register_scope_sanitize_start_guardrail( + &mut self, + name: &str, + priority: i32, + callback: EventSanitizeFn, + ) -> Result<()> { + let qualified_name = self.qualify_name(name); + register_scope_sanitize_start_guardrail(&qualified_name, priority, callback).map_err( + |err| PluginError::RegistrationFailed(format!("scope-start sanitizer: {err}")), + )?; + let name_owned = qualified_name; + self.registrations.push(PluginRegistration::new( + "plugin", + name_owned.clone(), + Box::new(move || { + deregister_scope_sanitize_start_guardrail(&name_owned) + .map(|_| ()) + .map_err(|err| { + PluginError::RegistrationFailed(format!( + "scope-start sanitizer deregistration failed: {err}" + )) + }) + }), + )); + Ok(()) + } + + /// Registers a scope-end event sanitizer and records its rollback closure. + pub fn register_scope_sanitize_end_guardrail( + &mut self, + name: &str, + priority: i32, + callback: EventSanitizeFn, + ) -> Result<()> { + let qualified_name = self.qualify_name(name); + register_scope_sanitize_end_guardrail(&qualified_name, priority, callback).map_err( + |err| PluginError::RegistrationFailed(format!("scope-end sanitizer: {err}")), + )?; + let name_owned = qualified_name; + self.registrations.push(PluginRegistration::new( + "plugin", + name_owned.clone(), + Box::new(move || { + deregister_scope_sanitize_end_guardrail(&name_owned) + .map(|_| ()) + .map_err(|err| { + PluginError::RegistrationFailed(format!( + "scope-end sanitizer deregistration failed: {err}" + )) + }) + }), + )); + Ok(()) + } + /// Registers an LLM request intercept and records its rollback closure. pub fn register_llm_request_intercept( &mut self, diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index c14f48f22..cfc409bf0 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -15,15 +15,15 @@ use std::task::{Context, Poll}; use chrono::{DateTime, Utc}; use libloading::{Library, Symbol}; use nemo_relay_plugin::{ - NEMO_RELAY_NATIVE_ABI_VERSION, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, - NemoRelayNativeHostApiV1, NemoRelayNativeJsonCb, NemoRelayNativeLlmConditionalCb, - NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCb, - NemoRelayNativeLlmRequestInterceptCb, NemoRelayNativeLlmStreamExecutionCb, - NemoRelayNativeLlmStreamV1, NemoRelayNativePluginContext, NemoRelayNativePluginEntry, - NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, NemoRelayNativeScopeStack, - NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, NemoRelayNativeString, - NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, NemoRelayNativeToolJsonCb, - NemoRelayNativeWithScopeStackCb, NemoRelayStatus, + NEMO_RELAY_NATIVE_ABI_VERSION, NemoRelayNativeEventSanitizeCb, + NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, + NemoRelayNativeJsonCb, NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, + NemoRelayNativeLlmRequestCb, NemoRelayNativeLlmRequestInterceptCb, + NemoRelayNativeLlmStreamExecutionCb, NemoRelayNativeLlmStreamV1, NemoRelayNativePluginContext, + NemoRelayNativePluginEntry, NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, + NemoRelayNativeScopeStack, NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, + NemoRelayNativeString, NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, + NemoRelayNativeToolJsonCb, NemoRelayNativeWithScopeStackCb, NemoRelayStatus, }; use semver::{Version, VersionReq}; use serde_json::{Map, Value as Json}; @@ -31,13 +31,13 @@ use sha2::{Digest, Sha256}; use tokio::runtime::Runtime; use tokio_stream::{Stream, StreamExt}; -use crate::api::event::Event; +use crate::api::event::{Event, EventSanitizeFields}; use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; use crate::api::runtime::{ - EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, LlmJsonStream, - LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn, - LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, - ToolInterceptFn, ToolSanitizeFn, + EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, + LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, + LlmStreamExecutionFn, LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionFn, + ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use crate::api::runtime::{ ScopeStackHandle, ThreadScopeStackBinding, capture_thread_scope_stack, create_scope_stack, @@ -597,6 +597,12 @@ fn native_host_api() -> *const NemoRelayNativeHostApiV1 { scope_stack_binding_free: native_scope_stack_binding_free, scope_stack_active: native_scope_stack_active, scope_stack_with_current: native_scope_stack_with_current, + plugin_context_register_mark_sanitize_guardrail: + native_plugin_context_register_mark_sanitize_guardrail, + plugin_context_register_scope_sanitize_start_guardrail: + native_plugin_context_register_scope_sanitize_start_guardrail, + plugin_context_register_scope_sanitize_end_guardrail: + native_plugin_context_register_scope_sanitize_end_guardrail, }) as *const _ } @@ -1401,6 +1407,52 @@ unsafe extern "C" fn native_plugin_context_register_llm_stream_execution_interce } } +macro_rules! native_event_sanitize_context_register { + ($fn_name:ident, $ctx_method:ident) => { + unsafe extern "C" fn $fn_name( + ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeEventSanitizeCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus { + clear_native_last_error(); + let host_ctx = match host_ctx_mut(ctx) { + Ok(ctx) => ctx, + Err(status) => return status, + }; + let instance = host_ctx.instance.clone(); + let ctx = unsafe { &mut *host_ctx.ctx }; + let name = match read_name(name) { + Ok(name) => name, + Err(status) => return status, + }; + match ctx.$ctx_method( + &name, + priority, + wrap_event_sanitize_fn(instance, cb, user_data, free_fn), + ) { + Ok(()) => NemoRelayStatus::Ok, + Err(err) => status_from_plugin_error(err), + } + } + }; +} + +native_event_sanitize_context_register!( + native_plugin_context_register_mark_sanitize_guardrail, + register_mark_sanitize_guardrail +); +native_event_sanitize_context_register!( + native_plugin_context_register_scope_sanitize_start_guardrail, + register_scope_sanitize_start_guardrail +); +native_event_sanitize_context_register!( + native_plugin_context_register_scope_sanitize_end_guardrail, + register_scope_sanitize_end_guardrail +); + fn wrap_event_subscriber( instance: Arc, cb: NemoRelayNativeEventSubscriberCb, @@ -1420,6 +1472,62 @@ fn wrap_event_subscriber( }) } +fn wrap_event_sanitize_fn( + instance: Arc, + cb: NemoRelayNativeEventSanitizeCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> EventSanitizeFn { + let user_data = make_user_data(instance, user_data, free_fn); + Arc::new(move |event, fields| { + let fallback = fields.clone(); + call_event_sanitize_callback(cb, user_data.ptr, event, &fields).unwrap_or(fallback) + }) +} + +fn call_event_sanitize_callback( + cb: NemoRelayNativeEventSanitizeCb, + user_data: *mut c_void, + event: &Event, + fields: &EventSanitizeFields, +) -> FlowResult { + clear_native_last_error(); + let event_json = serde_json::to_value(event) + .map_err(|err| FlowError::Internal(format!("failed to encode native event: {err}")))?; + let fields_json = serde_json::to_value(fields).map_err(|err| { + FlowError::Internal(format!("failed to encode native event fields: {err}")) + })?; + let event = native_string_from_json(&event_json) + .ok_or_else(|| FlowError::Internal("failed to allocate native event".into()))?; + let fields = match native_string_from_json(&fields_json) { + Some(fields) => fields, + None => { + unsafe { native_string_free(event) }; + return Err(FlowError::Internal( + "failed to allocate native event fields".into(), + )); + } + }; + let mut out = ptr::null_mut(); + let status = unsafe { cb(user_data, event, fields, &mut out) }; + unsafe { + native_string_free(event); + native_string_free(fields); + } + if status != NemoRelayStatus::Ok { + if !out.is_null() { + unsafe { native_string_free(out) }; + } + return Err(flow_error_from_status( + status, + "native event sanitizer failed", + )); + } + let value = take_json_from_native_string(out, "native event sanitizer returned null")?; + serde_json::from_value(value) + .map_err(|err| FlowError::Internal(format!("invalid event sanitize fields: {err}"))) +} + fn wrap_tool_json_fn( instance: Arc, cb: NemoRelayNativeToolJsonCb, diff --git a/crates/core/src/plugin/dynamic/worker.rs b/crates/core/src/plugin/dynamic/worker.rs index 5f7fda25f..797783584 100644 --- a/crates/core/src/plugin/dynamic/worker.rs +++ b/crates/core/src/plugin/dynamic/worker.rs @@ -51,7 +51,7 @@ use tokio_stream::wrappers::UnixListenerStream; #[cfg(unix)] use tower::service_fn; -use crate::api::event::Event; +use crate::api::event::{Event, EventSanitizeFields}; use crate::api::llm::LlmRequest; use crate::api::runtime::{ LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionNextFn, @@ -812,6 +812,29 @@ impl WorkerPluginInstance { }), )?; } + RegistrationSurface::MarkSanitizeGuardrail + | RegistrationSurface::ScopeSanitizeStartGuardrail + | RegistrationSurface::ScopeSanitizeEndGuardrail => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + let callback = Arc::new(move |event: &Event, fields: EventSanitizeFields| { + instance + .invoke_event_sanitize(&callback_name, surface, event) + .unwrap_or(fields) + }); + match surface { + RegistrationSurface::MarkSanitizeGuardrail => { + ctx.register_mark_sanitize_guardrail(&name, priority, callback)? + } + RegistrationSurface::ScopeSanitizeStartGuardrail => { + ctx.register_scope_sanitize_start_guardrail(&name, priority, callback)? + } + RegistrationSurface::ScopeSanitizeEndGuardrail => { + ctx.register_scope_sanitize_end_guardrail(&name, priority, callback)? + } + _ => unreachable!(), + } + } RegistrationSurface::ToolSanitizeRequestGuardrail => { let instance = Arc::new(self.clone_for_callback()); let callback_name = name.clone(); @@ -1143,6 +1166,26 @@ impl WorkerPluginCallback { } } + fn invoke_event_sanitize( + &self, + registration_name: &str, + surface: RegistrationSurface, + event: &Event, + ) -> FlowResult { + let request = self.base_request( + registration_name, + surface, + None, + Some(invoke_request_payload_event(event)), + ); + let value = json_from_invoke_response(self.invoke_blocking(request)?)?; + serde_json::from_value(value).map_err(|err| { + FlowError::Internal(format!( + "worker returned invalid event sanitize fields: {err}" + )) + }) + } + fn invoke_tool_json( &self, registration_name: &str, diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index 239228c83..f06e25f0d 100644 --- a/crates/core/src/stream.rs +++ b/crates/core/src/stream.rs @@ -37,6 +37,7 @@ use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; use crate::api::runtime::{EventSubscriberFn, ScopeStackHandle, current_scope_stack}; use crate::api::shared::metadata_with_otel_status; +use crate::api::shared::sanitize_event_with_scope_stack; use crate::codec::response::{AnnotatedLlmResponse, attach_estimated_cost_for_provider}; use crate::codec::traits::LlmResponseCodec; use crate::error::Result; @@ -226,7 +227,9 @@ impl LlmStreamWrapper { Err(_) => None, } }; - if let Some(event) = event_snapshot { + if let Some(event) = event_snapshot + && let Some(event) = sanitize_event_with_scope_stack(event, &self.scope_stack) + { NemoRelayContextState::emit_event(&event, &self.subscribers); } } @@ -253,7 +256,9 @@ impl LlmStreamWrapper { Err(_) => None, } }; - if let Some(event) = event_snapshot { + if let Some(event) = event_snapshot + && let Some(event) = sanitize_event_with_scope_stack(event, &self.scope_stack) + { NemoRelayContextState::emit_event(&event, &self.subscribers); } } diff --git a/crates/core/tests/fixtures/native_plugin/src/lib.rs b/crates/core/tests/fixtures/native_plugin/src/lib.rs index 4886a6c96..8dd5cf3de 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -5,8 +5,8 @@ use std::ffi::c_void; use std::ptr; use nemo_relay_plugin::{ - CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, Json, LlmJsonStream, - LlmRequest, LlmRequestInterceptOutcome, NemoRelayNativeHostApiV1, + CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, EventSanitizeFields, + Json, LlmJsonStream, LlmRequest, LlmRequestInterceptOutcome, NemoRelayNativeHostApiV1, NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeString, NemoRelayStatus, NemoRelayNativeToolNextFn, NativePlugin, PendingMarkSpec, PluginContext, PluginRuntime, ScopeCategory, ScopeType, ToolExecutionInterceptOutcome, @@ -47,6 +47,19 @@ impl NativePlugin for FixtureNativePlugin { let runtime = runtime.clone(); move |event| subscriber_mark(&runtime, event) })?; + ctx.register_mark_sanitize_guardrail("fixture_mark_sanitize", 0, |_, fields| { + mark_event_fields(fields, "native_plugin_mark") + })?; + ctx.register_scope_sanitize_start_guardrail( + "fixture_scope_start_sanitize", + 0, + |_, fields| mark_event_fields(fields, "native_plugin_scope_start"), + )?; + ctx.register_scope_sanitize_end_guardrail( + "fixture_scope_end_sanitize", + 0, + |_, fields| mark_event_fields(fields, "native_plugin_scope_end"), + )?; ctx.register_tool_sanitize_request_guardrail( "fixture_tool_sanitize_request", @@ -178,6 +191,13 @@ impl NativePlugin for FixtureNativePlugin { } } +fn mark_event_fields(mut fields: EventSanitizeFields, marker: &str) -> EventSanitizeFields { + let mut metadata = fields.metadata.unwrap_or_else(|| json!({})); + metadata[marker] = json!(true); + fields.metadata = Some(metadata); + fields +} + fn subscriber_mark(runtime: &PluginRuntime, event: &Event) { if event.name() == "native-plugin-test-outer" && event.scope_category() == Some(ScopeCategory::Start) diff --git a/crates/core/tests/fixtures/worker_plugin/src/main.rs b/crates/core/tests/fixtures/worker_plugin/src/main.rs index 4f8a17ead..44969c57a 100644 --- a/crates/core/tests/fixtures/worker_plugin/src/main.rs +++ b/crates/core/tests/fixtures/worker_plugin/src/main.rs @@ -6,7 +6,7 @@ use nemo_relay_worker::{ ToolExecutionInterceptOutcome, WorkerSdkError, serve_plugin, }; use nemo_relay_worker::{ - ConfigDiagnostic, DiagnosticLevel, Json, LlmRequest, PendingMarkSpec, + ConfigDiagnostic, DiagnosticLevel, EventSanitizeFields, Json, LlmRequest, PendingMarkSpec, }; use serde_json::json; @@ -61,6 +61,23 @@ impl WorkerPlugin for FixtureWorkerPlugin { let runtime = ctx .runtime() .ok_or_else(|| WorkerSdkError::Callback("runtime handle missing".into()))?; + ctx.register_mark_sanitize_guardrail("fixture_mark_sanitize", 0, |_, fields| { + mark_event_fields(fields, "worker_plugin_mark") + }); + ctx.register_mark_sanitize_guardrail("fixture_mark_sanitize_data", 1, |_, mut fields| { + fields.data = Some(json!({"worker_plugin_mark_data": true})); + fields + }); + ctx.register_scope_sanitize_start_guardrail( + "fixture_scope_start_sanitize", + 0, + |_, fields| mark_event_fields(fields, "worker_plugin_scope_start"), + ); + ctx.register_scope_sanitize_end_guardrail( + "fixture_scope_end_sanitize", + 0, + |_, fields| mark_event_fields(fields, "worker_plugin_scope_end"), + ); register_fixture_subscriber(ctx, runtime.clone()); register_fixture_tool_hooks( ctx, @@ -81,6 +98,13 @@ fn fixture_flag(config: &Json, key: &str) -> bool { config.get(key).and_then(Json::as_bool).unwrap_or(false) } +fn mark_event_fields(mut fields: EventSanitizeFields, marker: &str) -> EventSanitizeFields { + let mut metadata = fields.metadata.unwrap_or_else(|| json!({})); + metadata[marker] = json!(true); + fields.metadata = Some(metadata); + fields +} + fn register_fixture_subscriber(ctx: &mut PluginContext, runtime: nemo_relay_worker::PluginRuntime) { ctx.register_subscriber("fixture_subscriber", move |event| { if event.name() == "worker-plugin-test-outer" { diff --git a/crates/core/tests/integration/api_surface_tests.rs b/crates/core/tests/integration/api_surface_tests.rs index f35dc2137..fdac6e00e 100644 --- a/crates/core/tests/integration/api_surface_tests.rs +++ b/crates/core/tests/integration/api_surface_tests.rs @@ -10,7 +10,7 @@ use std::sync::{Arc, Mutex}; use chrono::{DateTime, TimeDelta, Utc}; use futures::StreamExt; -use nemo_relay::api::event::{Event, ScopeCategory}; +use nemo_relay::api::event::{CategoryProfile, Event, ScopeCategory}; use nemo_relay::api::llm::{LlmAttributes, LlmRequest}; use nemo_relay::api::llm::{ LlmCallExecuteParams, LlmCallParams, LlmStreamCallExecuteParams, llm_call, llm_call_end, @@ -20,18 +20,22 @@ use nemo_relay::api::registry::{ deregister_llm_conditional_execution_guardrail, deregister_llm_execution_intercept, deregister_llm_request_intercept, deregister_llm_sanitize_request_guardrail, deregister_llm_sanitize_response_guardrail, deregister_llm_stream_execution_intercept, - deregister_tool_conditional_execution_guardrail, deregister_tool_execution_intercept, - deregister_tool_request_intercept, deregister_tool_sanitize_request_guardrail, - deregister_tool_sanitize_response_guardrail, register_llm_conditional_execution_guardrail, - register_llm_execution_intercept, register_llm_request_intercept, - register_llm_sanitize_request_guardrail, register_llm_sanitize_response_guardrail, - register_llm_stream_execution_intercept, register_tool_conditional_execution_guardrail, + deregister_mark_sanitize_guardrail, deregister_scope_sanitize_end_guardrail, + deregister_scope_sanitize_start_guardrail, deregister_tool_conditional_execution_guardrail, + deregister_tool_execution_intercept, deregister_tool_request_intercept, + deregister_tool_sanitize_request_guardrail, deregister_tool_sanitize_response_guardrail, + register_llm_conditional_execution_guardrail, register_llm_execution_intercept, + register_llm_request_intercept, register_llm_sanitize_request_guardrail, + register_llm_sanitize_response_guardrail, register_llm_stream_execution_intercept, + register_mark_sanitize_guardrail, register_scope_sanitize_end_guardrail, + register_scope_sanitize_start_guardrail, register_tool_conditional_execution_guardrail, register_tool_execution_intercept, register_tool_request_intercept, register_tool_sanitize_request_guardrail, register_tool_sanitize_response_guardrail, scope_deregister_llm_conditional_execution_guardrail, scope_deregister_llm_execution_intercept, scope_deregister_llm_request_intercept, scope_deregister_llm_sanitize_request_guardrail, scope_deregister_llm_sanitize_response_guardrail, - scope_deregister_llm_stream_execution_intercept, + scope_deregister_llm_stream_execution_intercept, scope_deregister_mark_sanitize_guardrail, + scope_deregister_scope_sanitize_end_guardrail, scope_deregister_scope_sanitize_start_guardrail, scope_deregister_tool_conditional_execution_guardrail, scope_deregister_tool_execution_intercept, scope_deregister_tool_request_intercept, scope_deregister_tool_sanitize_request_guardrail, @@ -39,6 +43,8 @@ use nemo_relay::api::registry::{ scope_register_llm_conditional_execution_guardrail, scope_register_llm_execution_intercept, scope_register_llm_request_intercept, scope_register_llm_sanitize_request_guardrail, scope_register_llm_sanitize_response_guardrail, scope_register_llm_stream_execution_intercept, + scope_register_mark_sanitize_guardrail, scope_register_scope_sanitize_end_guardrail, + scope_register_scope_sanitize_start_guardrail, scope_register_tool_conditional_execution_guardrail, scope_register_tool_execution_intercept, scope_register_tool_request_intercept, scope_register_tool_sanitize_request_guardrail, scope_register_tool_sanitize_response_guardrail, @@ -71,6 +77,209 @@ fn reset_global() { *state = NemoRelayContextState::new(); } +#[test] +fn event_sanitizers_rewrite_only_observability_fields_in_priority_order() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + let events = capture_events("event-sanitizer-fields"); + + register_scope_sanitize_start_guardrail( + "scope-start-late", + 20, + Arc::new(|event, mut fields| { + assert_eq!(event.data().unwrap()["order"], json!(["early"])); + fields.data = Some(json!({"order": ["early", "late"]})); + fields + }), + ) + .unwrap(); + register_scope_sanitize_start_guardrail( + "scope-start-early", + 10, + Arc::new(|_, mut fields| { + fields.data = Some(json!({"order": ["early"]})); + fields.metadata = Some(json!({"redacted": true})); + fields.category_profile = Some(CategoryProfile::builder().subtype("sanitized").build()); + fields + }), + ) + .unwrap(); + + let handle = push_scope( + nemo_relay::api::scope::PushScopeParams::builder() + .name("sanitized-scope") + .scope_type(ScopeType::Custom) + .input(json!({"secret": "raw"})) + .metadata(json!({"secret": "raw"})) + .build(), + ) + .unwrap(); + let snapshot = captured_events_snapshot(&events); + let start = snapshot + .iter() + .find(|event| event.name() == "sanitized-scope") + .unwrap(); + assert_eq!(start.data().unwrap(), &json!({"order": ["early", "late"]})); + assert_eq!(start.metadata().unwrap(), &json!({"redacted": true})); + assert_eq!( + start.category_profile().unwrap().subtype.as_deref(), + Some("sanitized") + ); + assert_eq!(start.uuid(), handle.uuid); + assert_eq!(start.category().unwrap().as_str(), "custom"); + + pop_scope( + nemo_relay::api::scope::PopScopeParams::builder() + .handle_uuid(&handle.uuid) + .build(), + ) + .unwrap(); + deregister_scope_sanitize_start_guardrail("scope-start-early").unwrap(); + deregister_scope_sanitize_start_guardrail("scope-start-late").unwrap(); + deregister_subscriber("event-sanitizer-fields").unwrap(); +} + +#[test] +fn mark_and_scope_local_sanitizers_cover_marks_and_tool_scopes() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + let events = capture_events("event-sanitizer-kinds"); + register_mark_sanitize_guardrail( + "mark-global", + 20, + Arc::new(|_, mut fields| { + fields.data = Some(json!({"mark": "global"})); + fields + }), + ) + .unwrap(); + register_scope_sanitize_end_guardrail( + "scope-end-global", + 20, + Arc::new(|_, mut fields| { + fields.metadata = Some(json!({"scope_end": true})); + fields + }), + ) + .unwrap(); + + let owner = push_scope( + nemo_relay::api::scope::PushScopeParams::builder() + .name("owner") + .scope_type(ScopeType::Agent) + .build(), + ) + .unwrap(); + scope_register_mark_sanitize_guardrail( + &owner.uuid, + "mark-local", + 10, + Arc::new(|_, mut fields| { + fields.data = Some(json!({"mark": "local"})); + fields + }), + ) + .unwrap(); + scope_register_scope_sanitize_start_guardrail( + &owner.uuid, + "scope-start-local", + 10, + Arc::new(|_, mut fields| { + fields.metadata = Some(json!({"scope_start": true})); + fields + }), + ) + .unwrap(); + scope_register_scope_sanitize_end_guardrail( + &owner.uuid, + "scope-end-local", + 10, + Arc::new(|_, mut fields| { + fields.data = Some(json!({"scope_end": "local"})); + fields + }), + ) + .unwrap(); + + event( + nemo_relay::api::scope::EmitMarkEventParams::builder() + .name("sanitized-mark") + .data(json!({"secret": "raw"})) + .build(), + ) + .unwrap(); + event( + nemo_relay::api::scope::EmitMarkEventParams::builder() + .name("sanitized-empty-mark") + .build(), + ) + .unwrap(); + let tool = tool_call( + nemo_relay::api::tool::ToolCallParams::builder() + .name("sanitized-tool") + .args(json!({"input": true})) + .build(), + ) + .unwrap(); + tool_call_end( + nemo_relay::api::tool::ToolCallEndParams::builder() + .handle(&tool) + .result(json!({"output": true})) + .build(), + ) + .unwrap(); + + let snapshot = captured_events_snapshot(&events); + let mark = snapshot + .iter() + .find(|event| event.name() == "sanitized-mark") + .unwrap(); + assert_eq!(mark.data().unwrap(), &json!({"mark": "global"})); + let empty_mark = snapshot + .iter() + .find(|event| event.name() == "sanitized-empty-mark") + .unwrap(); + assert_eq!(empty_mark.data().unwrap(), &json!({"mark": "global"})); + let tool_start = snapshot + .iter() + .find(|event| { + event.name() == "sanitized-tool" && event.scope_category() == Some(ScopeCategory::Start) + }) + .unwrap(); + assert_eq!(tool_start.category().unwrap().as_str(), "tool"); + assert_eq!( + tool_start.metadata().unwrap(), + &json!({"scope_start": true}) + ); + let tool_end = snapshot + .iter() + .find(|event| { + event.name() == "sanitized-tool" && event.scope_category() == Some(ScopeCategory::End) + }) + .unwrap(); + assert_eq!(tool_end.data().unwrap(), &json!({"scope_end": "local"})); + assert_eq!(tool_end.metadata().unwrap(), &json!({"scope_end": true})); + + pop_scope( + nemo_relay::api::scope::PopScopeParams::builder() + .handle_uuid(&owner.uuid) + .build(), + ) + .unwrap(); + let snapshot = captured_events_snapshot(&events); + let owner_end = snapshot + .iter() + .find(|event| event.name() == "owner" && event.scope_category() == Some(ScopeCategory::End)) + .unwrap(); + assert_eq!(owner_end.data().unwrap(), &json!({"scope_end": "local"})); + assert_eq!(owner_end.metadata().unwrap(), &json!({"scope_end": true})); + deregister_mark_sanitize_guardrail("mark-global").unwrap(); + deregister_scope_sanitize_end_guardrail("scope-end-global").unwrap(); + deregister_subscriber("event-sanitizer-kinds").unwrap(); +} + fn setup_isolated_thread() { let stack = create_scope_stack(); set_thread_scope_stack(stack); @@ -353,6 +562,29 @@ fn test_global_registry_and_subscriber_wrappers_cover_success_and_duplicates() { reset_global(); setup_isolated_thread(); + register_mark_sanitize_guardrail("mark-sanitize", 1, Arc::new(|_, fields| fields)).unwrap(); + expect_already_exists( + register_mark_sanitize_guardrail("mark-sanitize", 1, Arc::new(|_, fields| fields)) + .unwrap_err(), + "mark-sanitize", + ); + assert!(deregister_mark_sanitize_guardrail("mark-sanitize").unwrap()); + assert!(!deregister_mark_sanitize_guardrail("mark-sanitize").unwrap()); + + register_scope_sanitize_start_guardrail( + "scope-start-sanitize", + 1, + Arc::new(|_, fields| fields), + ) + .unwrap(); + assert!(deregister_scope_sanitize_start_guardrail("scope-start-sanitize").unwrap()); + assert!(!deregister_scope_sanitize_start_guardrail("scope-start-sanitize").unwrap()); + + register_scope_sanitize_end_guardrail("scope-end-sanitize", 1, Arc::new(|_, fields| fields)) + .unwrap(); + assert!(deregister_scope_sanitize_end_guardrail("scope-end-sanitize").unwrap()); + assert!(!deregister_scope_sanitize_end_guardrail("scope-end-sanitize").unwrap()); + register_tool_sanitize_request_guardrail( "tool-sanitize-request", 1, @@ -510,6 +742,56 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss ) .unwrap(); + scope_register_mark_sanitize_guardrail( + &scope.uuid, + "mark-sanitize", + 1, + Arc::new(|_, fields| fields), + ) + .unwrap(); + expect_already_exists( + scope_register_mark_sanitize_guardrail( + &scope.uuid, + "mark-sanitize", + 1, + Arc::new(|_, fields| fields), + ) + .unwrap_err(), + "mark-sanitize", + ); + assert!(scope_deregister_mark_sanitize_guardrail(&scope.uuid, "mark-sanitize").unwrap()); + assert!(!scope_deregister_mark_sanitize_guardrail(&scope.uuid, "mark-sanitize").unwrap()); + + scope_register_scope_sanitize_start_guardrail( + &scope.uuid, + "scope-start-sanitize", + 1, + Arc::new(|_, fields| fields), + ) + .unwrap(); + assert!( + scope_deregister_scope_sanitize_start_guardrail(&scope.uuid, "scope-start-sanitize") + .unwrap() + ); + assert!( + !scope_deregister_scope_sanitize_start_guardrail(&scope.uuid, "scope-start-sanitize") + .unwrap() + ); + + scope_register_scope_sanitize_end_guardrail( + &scope.uuid, + "scope-end-sanitize", + 1, + Arc::new(|_, fields| fields), + ) + .unwrap(); + assert!( + scope_deregister_scope_sanitize_end_guardrail(&scope.uuid, "scope-end-sanitize").unwrap() + ); + assert!( + !scope_deregister_scope_sanitize_end_guardrail(&scope.uuid, "scope-end-sanitize").unwrap() + ); + scope_register_tool_sanitize_request_guardrail( &scope.uuid, "tool-sanitize-request", @@ -664,6 +946,36 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss ) .unwrap(); + expect_not_found( + scope_register_mark_sanitize_guardrail( + &scope.uuid, + "missing-mark-sanitize", + 1, + Arc::new(|_, fields| fields), + ) + .unwrap_err(), + "scope", + ); + expect_not_found( + scope_register_scope_sanitize_start_guardrail( + &scope.uuid, + "missing-scope-start-sanitize", + 1, + Arc::new(|_, fields| fields), + ) + .unwrap_err(), + "scope", + ); + expect_not_found( + scope_register_scope_sanitize_end_guardrail( + &scope.uuid, + "missing-scope-end-sanitize", + 1, + Arc::new(|_, fields| fields), + ) + .unwrap_err(), + "scope", + ); expect_not_found( scope_register_tool_sanitize_request_guardrail( &scope.uuid, @@ -733,6 +1045,30 @@ async fn test_tool_api_emits_sanitized_events_and_covers_error_paths() { }), ) .unwrap(); + register_scope_sanitize_start_guardrail( + "tool-generic-sanitize-start", + 1, + Arc::new(|event, mut fields| { + if event.name() == "tool-api" { + assert_eq!(event.input().unwrap()["sanitized_request"], true); + fields.metadata = Some(json!({"generic_start": true})); + } + fields + }), + ) + .unwrap(); + register_scope_sanitize_end_guardrail( + "tool-generic-sanitize-end", + 1, + Arc::new(|event, mut fields| { + if event.name() == "tool-api" { + assert_eq!(event.output().unwrap()["sanitized_response"], true); + fields.metadata = Some(json!({"generic_end": true})); + } + fields + }), + ) + .unwrap(); let handle = tool_call( nemo_relay::api::tool::ToolCallParams::builder() @@ -764,6 +1100,7 @@ async fn test_tool_api_emits_sanitized_events_and_covers_error_paths() { json!(true) ); assert_eq!(captured[0].tool_call_id(), Some("tool-call-id")); + assert_eq!(captured[0].metadata().unwrap()["generic_start"], true); assert_eq!(captured[1].kind(), "scope"); assert_eq!(captured[1].scope_category(), Some(ScopeCategory::End)); assert_eq!(captured[1].category().unwrap().as_str(), "tool"); @@ -772,8 +1109,11 @@ async fn test_tool_api_emits_sanitized_events_and_covers_error_paths() { json!(true) ); assert_eq!(captured[1].tool_call_id(), Some("tool-call-id")); + assert_eq!(captured[1].metadata().unwrap()["generic_end"], true); deregister_tool_sanitize_request_guardrail("tool-sanitize-request").unwrap(); deregister_tool_sanitize_response_guardrail("tool-sanitize-response").unwrap(); + deregister_scope_sanitize_start_guardrail("tool-generic-sanitize-start").unwrap(); + deregister_scope_sanitize_end_guardrail("tool-generic-sanitize-end").unwrap(); register_tool_request_intercept( "tool-request", @@ -1037,6 +1377,16 @@ async fn test_llm_stream_chunk_marks_track_successful_chunks() { setup_isolated_thread(); let events = capture_events("llm-stream-chunk-mark-events"); + register_mark_sanitize_guardrail( + "stream-mark-sanitize", + 1, + Arc::new(|event, mut fields| { + assert_eq!(event.name(), "llm.chunk"); + fields.metadata = Some(json!({"sanitized": true})); + fields + }), + ) + .unwrap(); let raw_chunks = vec![ json!({ "object": "chat.completion.chunk", @@ -1094,6 +1444,7 @@ async fn test_llm_stream_chunk_marks_track_successful_chunks() { assert_eq!(mark.parent_uuid(), Some(llm_uuid)); assert_eq!(mark.category(), None); assert_eq!(mark.scope_type(), None); + assert_eq!(mark.metadata().unwrap(), &json!({"sanitized": true})); } assert_eq!( captured[1].data().unwrap(), @@ -1109,6 +1460,7 @@ async fn test_llm_stream_chunk_marks_track_successful_chunks() { &json!({"chunk_index": 1, "provider": "unknown"}) ); + deregister_mark_sanitize_guardrail("stream-mark-sanitize").unwrap(); deregister_subscriber("llm-stream-chunk-mark-events").unwrap(); } diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index b978195bf..98346fc13 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -26,12 +26,13 @@ use nemo_relay::api::registry::{ deregister_llm_conditional_execution_guardrail, deregister_llm_execution_intercept, deregister_llm_request_intercept, deregister_llm_sanitize_request_guardrail, deregister_llm_sanitize_response_guardrail, deregister_llm_stream_execution_intercept, - deregister_tool_conditional_execution_guardrail, deregister_tool_execution_intercept, - deregister_tool_request_intercept, deregister_tool_sanitize_request_guardrail, - deregister_tool_sanitize_response_guardrail, register_llm_conditional_execution_guardrail, - register_llm_execution_intercept, register_llm_request_intercept, - register_llm_sanitize_request_guardrail, register_llm_sanitize_response_guardrail, - register_llm_stream_execution_intercept, register_tool_conditional_execution_guardrail, + deregister_mark_sanitize_guardrail, deregister_tool_conditional_execution_guardrail, + deregister_tool_execution_intercept, deregister_tool_request_intercept, + deregister_tool_sanitize_request_guardrail, deregister_tool_sanitize_response_guardrail, + register_llm_conditional_execution_guardrail, register_llm_execution_intercept, + register_llm_request_intercept, register_llm_sanitize_request_guardrail, + register_llm_sanitize_response_guardrail, register_llm_stream_execution_intercept, + register_mark_sanitize_guardrail, register_tool_conditional_execution_guardrail, register_tool_execution_intercept, register_tool_request_intercept, register_tool_sanitize_request_guardrail, register_tool_sanitize_response_guardrail, scope_register_llm_conditional_execution_guardrail, scope_register_llm_execution_intercept, @@ -670,6 +671,17 @@ async fn test_tool_execution_outcome_marks_follow_end_with_tool_parentage() { Arc::new(move |event| captured.lock().unwrap().push(event.clone())), ) .unwrap(); + register_mark_sanitize_guardrail( + "tool_pending_mark_sanitizer", + 1, + Arc::new(|_, mut fields| { + let mut metadata = fields.metadata.unwrap_or_else(|| json!({})); + metadata["sanitized"] = json!(true); + fields.metadata = Some(metadata); + fields + }), + ) + .unwrap(); let mut plugin_ctx = PluginRegistrationContext::new(); plugin_ctx @@ -759,6 +771,8 @@ async fn test_tool_execution_outcome_marks_follow_end_with_tool_parentage() { .iter() .position(|event| event.name() == "tool.mark.outer") .unwrap(); + assert_eq!(captured[inner_index].metadata().unwrap()["sanitized"], true); + assert_eq!(captured[outer_index].metadata().unwrap()["sanitized"], true); assert!(start_index < end_index); assert!(end_index < inner_index); assert!(inner_index < outer_index); @@ -786,6 +800,7 @@ async fn test_tool_execution_outcome_marks_follow_end_with_tool_parentage() { deregister_tool_execution_intercept("passthrough_between_outcomes").unwrap(); let mut registrations = plugin_ctx.into_registrations(); rollback_registrations(&mut registrations); + deregister_mark_sanitize_guardrail("tool_pending_mark_sanitizer").unwrap(); deregister_subscriber("tool_outcome_mark_observer").unwrap(); } @@ -3104,6 +3119,15 @@ async fn test_managed_llm_emits_pending_marks_under_started_scope() { Arc::new(move |event: &Event| captured.lock().unwrap().push(event.clone())), ) .unwrap(); + register_mark_sanitize_guardrail( + "llm_pending_mark_sanitizer", + 1, + Arc::new(|event, mut fields| { + fields.metadata = Some(json!({"sanitized_mark": event.name()})); + fields + }), + ) + .unwrap(); register_llm_request_intercept( "pending_managed", @@ -3191,8 +3215,17 @@ async fn test_managed_llm_emits_pending_marks_under_started_scope() { assert_eq!(mark.timestamp(), second_mark.timestamp()); assert!(end.timestamp() >= mark.timestamp()); assert_eq!(mark.data().unwrap()["saved_tokens"], 12); + assert_eq!( + mark.metadata().unwrap()["sanitized_mark"], + "request.optimized" + ); + assert_eq!( + second_mark.metadata().unwrap()["sanitized_mark"], + "request.optimized.second" + ); deregister_llm_request_intercept("pending_managed").unwrap(); + deregister_mark_sanitize_guardrail("llm_pending_mark_sanitizer").unwrap(); deregister_subscriber("pending_mark_observer").unwrap(); } diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index 85f6a4b69..73496a910 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -162,12 +162,38 @@ async fn sdk_cdylib_registers_tool_request_intercept() { let first_events = events.lock().unwrap().clone(); find_event(&first_events, "fixture.native.subscriber.mark", None); assert_parent(&first_events, "fixture.native.mark", None, Some(outer_uuid)); + assert_eq!( + find_event(&first_events, "fixture.native.mark", None) + .metadata() + .unwrap()["native_plugin_mark"], + true + ); assert_parent( &first_events, "fixture.native.scope", Some(ScopeCategory::Start), Some(outer_uuid), ); + assert_eq!( + find_event( + &first_events, + "fixture.native.scope", + Some(ScopeCategory::Start), + ) + .metadata() + .unwrap()["native_plugin_scope_start"], + true + ); + assert_eq!( + find_event( + &first_events, + "fixture.native.scope", + Some(ScopeCategory::End), + ) + .metadata() + .unwrap()["native_plugin_scope_end"], + true + ); assert_not_parent( &first_events, "fixture.native.isolated.mark", @@ -189,6 +215,10 @@ async fn sdk_cdylib_registers_tool_request_intercept() { tool_start.input().unwrap()["native_plugin_tool_sanitize_request"], true ); + assert_eq!( + tool_start.metadata().unwrap()["native_plugin_scope_start"], + true + ); let tool_end = find_event( &first_events, "native-fixture-tool", @@ -198,6 +228,10 @@ async fn sdk_cdylib_registers_tool_request_intercept() { tool_end.output().unwrap()["native_plugin_tool_sanitize_response"], true ); + assert_eq!( + tool_end.metadata().unwrap()["native_plugin_scope_end"], + true + ); assert!(tool_end.output().unwrap().get("pending_marks").is_none()); let tool_mark = find_event(&first_events, "fixture.native.tool_execution.mark", None); assert_eq!(tool_mark.parent_uuid(), Some(tool_start.uuid())); @@ -213,6 +247,7 @@ async fn sdk_cdylib_registers_tool_request_intercept() { ); assert_eq!(tool_mark.data().unwrap()["source"], "native_tool_execution"); assert_eq!(tool_mark.metadata().unwrap()["fixture"], true); + assert_eq!(tool_mark.metadata().unwrap()["native_plugin_mark"], true); assert!(tool_mark.timestamp() > tool_end.timestamp()); let tool_end_index = first_events .iter() diff --git a/crates/core/tests/integration/worker_plugin_tests.rs b/crates/core/tests/integration/worker_plugin_tests.rs index 842b6c191..d8b8e21cc 100644 --- a/crates/core/tests/integration/worker_plugin_tests.rs +++ b/crates/core/tests/integration/worker_plugin_tests.rs @@ -14,7 +14,9 @@ use nemo_relay::api::llm::{ llm_stream_call_execute, }; use nemo_relay::api::runtime::{TASK_SCOPE_STACK, create_scope_stack}; -use nemo_relay::api::scope::{PopScopeParams, PushScopeParams, ScopeType, pop_scope, push_scope}; +use nemo_relay::api::scope::{ + EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeType, event, pop_scope, push_scope, +}; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use nemo_relay::api::tool::{ToolCallExecuteParams, tool_call_execute, tool_request_intercepts}; use nemo_relay::codec::request::AnnotatedLlmRequest; @@ -102,12 +104,38 @@ async fn rust_worker_registers_and_invokes_all_current_surfaces() { None, Some(outer_uuid), ); + assert_eq!( + find_event(&captured_events, "fixture.worker.mark", None) + .metadata() + .unwrap()["worker_plugin_mark"], + true + ); assert_parent( &captured_events, "fixture.worker.scope", Some(ScopeCategory::Start), Some(outer_uuid), ); + assert_eq!( + find_event( + &captured_events, + "fixture.worker.scope", + Some(ScopeCategory::Start), + ) + .metadata() + .unwrap()["worker_plugin_scope_start"], + true + ); + assert_eq!( + find_event( + &captured_events, + "fixture.worker.scope", + Some(ScopeCategory::End), + ) + .metadata() + .unwrap()["worker_plugin_scope_end"], + true + ); assert_not_parent( &captured_events, "fixture.worker.isolated.scope", @@ -135,6 +163,10 @@ async fn rust_worker_registers_and_invokes_all_current_surfaces() { tool_start.input().unwrap()["worker_plugin_tool_sanitize_request"], true ); + assert_eq!( + tool_start.metadata().unwrap()["worker_plugin_scope_start"], + true + ); let tool_end = find_event( &captured_events, "worker-fixture-tool", @@ -144,9 +176,14 @@ async fn rust_worker_registers_and_invokes_all_current_surfaces() { tool_end.output().unwrap()["worker_plugin_tool_sanitize_response"], true ); + assert_eq!( + tool_end.metadata().unwrap()["worker_plugin_scope_end"], + true + ); let tool_mark = find_event(&captured_events, "fixture.worker.tool_execution.mark", None); assert_eq!(tool_mark.parent_uuid(), Some(tool_start.uuid())); assert!(tool_mark.timestamp() > tool_end.timestamp()); + assert_eq!(tool_mark.metadata().unwrap()["worker_plugin_mark"], true); let llm_execute_response = llm_call_execute( LlmCallExecuteParams::builder() @@ -189,10 +226,11 @@ async fn rust_worker_registers_and_invokes_all_current_surfaces() { let pending_mark = find_event(&captured_events, "fixture.worker.llm_request.mark", None); assert_eq!(pending_mark.parent_uuid(), Some(llm_start.uuid())); assert_eq!( - pending_mark.data().unwrap()["source"], - "worker_request_intercept" + pending_mark.data().unwrap()["worker_plugin_mark_data"], + true ); assert_eq!(pending_mark.metadata().unwrap()["fixture"], true); + assert_eq!(pending_mark.metadata().unwrap()["worker_plugin_mark"], true); let llm_end = find_event( &captured_events, "worker-fixture-llm-execute", @@ -241,6 +279,47 @@ async fn rust_worker_registers_and_invokes_all_current_surfaces() { loaded.clear(); } +#[tokio::test] +async fn worker_event_sanitizers_preserve_prior_field_changes() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let loaded = load_and_initialize_fixture(Map::new()).await; + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "worker_event_sanitizer_field_preservation", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .expect("test subscriber should register"); + + TASK_SCOPE_STACK + .scope(create_scope_stack(), async { + event( + EmitMarkEventParams::builder() + .name("worker-event-sanitizer-field-preservation") + .data(json!({"original_data": true})) + .metadata(json!({"original_metadata": true})) + .build(), + ) + .expect("mark event should emit"); + }) + .await; + + flush_subscribers().expect("worker fixture events should flush"); + let captured_events = events.lock().unwrap().clone(); + let mark = find_event( + &captured_events, + "worker-event-sanitizer-field-preservation", + None, + ); + assert_eq!(mark.data(), Some(&json!({"worker_plugin_mark_data": true}))); + assert_eq!(mark.metadata().unwrap()["worker_plugin_mark"], true); + assert_eq!(mark.metadata().unwrap()["original_metadata"], true); + + deregister_subscriber("worker_event_sanitizer_field_preservation") + .expect("test subscriber should deregister"); + loaded.clear(); +} + #[tokio::test] async fn host_cancellation_reaches_rust_worker_invocation() { let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; diff --git a/crates/core/tests/unit/dynamic_worker_tests.rs b/crates/core/tests/unit/dynamic_worker_tests.rs index cd8abf3b8..e43a9586b 100644 --- a/crates/core/tests/unit/dynamic_worker_tests.rs +++ b/crates/core/tests/unit/dynamic_worker_tests.rs @@ -376,7 +376,7 @@ async fn callback_helpers_cover_worker_response_edges() { error: None, })), }, - "llm_json_invalid" => InvokeResponse { + "llm_json_invalid" | "event_fields_invalid" => InvokeResponse { result: Some(InvokeResult::Json(JsonResult { value: Some(json_envelope(JSON_SCHEMA, &json!(null)).expect("json envelope")), error: None, @@ -436,6 +436,19 @@ async fn callback_helpers_cover_worker_response_edges() { .expect_err("unexpected subscriber result should fail"); assert!(error.to_string().contains("subscriber returned unexpected")); + let error = callback + .invoke_event_sanitize( + "event_fields_invalid", + RegistrationSurface::MarkSanitizeGuardrail, + &event, + ) + .expect_err("invalid event sanitizer fields should fail"); + assert!( + error + .to_string() + .contains("worker returned invalid event sanitize fields") + ); + let error = callback .invoke_llm_request_json( "llm_json_invalid", diff --git a/crates/core/tests/unit/plugin_tests.rs b/crates/core/tests/unit/plugin_tests.rs index 282c5c571..bafe55652 100644 --- a/crates/core/tests/unit/plugin_tests.rs +++ b/crates/core/tests/unit/plugin_tests.rs @@ -983,6 +983,20 @@ fn test_plugin_registration_context_supports_guardrail_helpers() { reset_global(); let mut ctx = PluginRegistrationContext::with_namespace("plugin::"); + ctx.register_mark_sanitize_guardrail("mark_sanitize", 1, Arc::new(|_, fields| fields)) + .unwrap(); + ctx.register_scope_sanitize_start_guardrail( + "scope_sanitize_start", + 1, + Arc::new(|_, fields| fields), + ) + .unwrap(); + ctx.register_scope_sanitize_end_guardrail( + "scope_sanitize_end", + 1, + Arc::new(|_, fields| fields), + ) + .unwrap(); ctx.register_tool_sanitize_request_guardrail( "tool_sanitize_request", 1, @@ -1057,6 +1071,24 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { reset_global(); let mut ctx = PluginRegistrationContext::with_namespace("duplicate::"); + ctx.register_mark_sanitize_guardrail("mark", 1, Arc::new(|_, fields| fields)) + .unwrap(); + expect_registration_failed( + ctx.register_mark_sanitize_guardrail("mark", 1, Arc::new(|_, fields| fields)), + "mark sanitizer:", + ); + ctx.register_scope_sanitize_start_guardrail("scope-start", 1, Arc::new(|_, fields| fields)) + .unwrap(); + expect_registration_failed( + ctx.register_scope_sanitize_start_guardrail("scope-start", 1, Arc::new(|_, fields| fields)), + "scope-start sanitizer:", + ); + ctx.register_scope_sanitize_end_guardrail("scope-end", 1, Arc::new(|_, fields| fields)) + .unwrap(); + expect_registration_failed( + ctx.register_scope_sanitize_end_guardrail("scope-end", 1, Arc::new(|_, fields| fields)), + "scope-end sanitizer:", + ); ctx.register_llm_request_intercept( "llm-request", 1, @@ -1246,6 +1278,20 @@ fn test_plugin_registration_context_maps_deregistration_errors() { reset_global(); let mut ctx = PluginRegistrationContext::with_namespace("teardown::"); + ctx.register_mark_sanitize_guardrail("mark-sanitize", 1, Arc::new(|_, fields| fields)) + .unwrap(); + ctx.register_scope_sanitize_start_guardrail( + "scope-sanitize-start", + 1, + Arc::new(|_, fields| fields), + ) + .unwrap(); + ctx.register_scope_sanitize_end_guardrail( + "scope-sanitize-end", + 1, + Arc::new(|_, fields| fields), + ) + .unwrap(); ctx.register_subscriber("subscriber", Arc::new(|_event| {})) .unwrap(); ctx.register_llm_request_intercept( @@ -1319,6 +1365,9 @@ fn test_plugin_registration_context_maps_deregistration_errors() { let mut registrations = ctx.into_registrations(); let expected_messages = [ + "mark sanitizer deregistration failed:", + "scope-start sanitizer deregistration failed:", + "scope-end sanitizer deregistration failed:", "subscriber deregistration failed:", "llm request intercept deregistration failed:", "tool sanitize request guardrail deregistration failed:", diff --git a/crates/core/tests/unit/types_tests.rs b/crates/core/tests/unit/types_tests.rs index 7f745935d..dc899e46c 100644 --- a/crates/core/tests/unit/types_tests.rs +++ b/crates/core/tests/unit/types_tests.rs @@ -10,9 +10,9 @@ use serde_json::{Map, json}; use uuid::{Uuid, Version}; use crate::api::event::{ - BaseEvent, CategoryProfile, DataSchema, Event, EventCategory, EventNormalizationExt, MarkEvent, - ScopeCategory, ScopeEvent, attributes_from_handle, llm_attributes_to_strings, - scope_attributes_to_strings, tool_attributes_to_strings, + BaseEvent, CategoryProfile, DataSchema, Event, EventCategory, EventNormalizationExt, + EventSanitizeFields, MarkEvent, ScopeCategory, ScopeEvent, attributes_from_handle, + llm_attributes_to_strings, scope_attributes_to_strings, tool_attributes_to_strings, }; use crate::api::llm::{LlmAttributes, LlmHandle, LlmRequest}; use crate::api::scope::{HandleAttributes, ScopeAttributes, ScopeHandle, ScopeType}; @@ -948,3 +948,45 @@ fn event_scope_accessors_return_none_for_unprofiled_and_wrong_phase_payloads() { assert_eq!(end_without_data.input(), None); assert_eq!(end_without_data.output(), None); } + +#[test] +fn event_sanitize_fields_snapshot_and_apply_cover_scope_and_mark_variants() { + let profile = CategoryProfile::builder().subtype("raw").build(); + let mut scope = Event::Scope(ScopeEvent::new( + BaseEvent::builder() + .name("scope") + .data(json!({"secret": "raw"})) + .metadata(json!({"meta": "raw"})) + .build(), + ScopeCategory::Start, + Vec::new(), + EventCategory::custom(), + Some(profile), + )); + assert_eq!(scope.sanitize_fields().data, Some(json!({"secret": "raw"}))); + scope.apply_sanitize_fields(EventSanitizeFields { + data: Some(json!({"secret": "clean"})), + category_profile: None, + metadata: None, + }); + assert_eq!(scope.data(), Some(&json!({"secret": "clean"}))); + assert!(scope.category_profile().is_none()); + assert!(scope.metadata().is_none()); + + let original_uuid = Uuid::now_v7(); + let mut mark = Event::Mark(MarkEvent::new( + BaseEvent::builder() + .name("mark") + .uuid(original_uuid) + .build(), + None, + None, + )); + mark.apply_sanitize_fields( + EventSanitizeFields::builder() + .metadata(json!({"added": true})) + .build(), + ); + assert_eq!(mark.uuid(), original_uuid); + assert_eq!(mark.metadata(), Some(&json!({"added": true}))); +} diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 3f96f9826..9f1dfd918 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -17,7 +17,7 @@ use std::sync::Mutex; pub use nemo_relay_types::Json; pub use nemo_relay_types::api::event::{ - CategoryProfile, Event, EventCategory, PendingMarkSpec, ScopeCategory, + CategoryProfile, Event, EventCategory, EventSanitizeFields, PendingMarkSpec, ScopeCategory, }; pub use nemo_relay_types::api::llm::{LlmAttributes, LlmRequest, LlmRequestInterceptOutcome}; pub use nemo_relay_types::api::scope::{HandleAttributes, ScopeAttributes, ScopeType}; @@ -200,6 +200,14 @@ pub type NemoRelayNativeEventSubscriberCb = unsafe extern "C" fn( event_json: *const NemoRelayNativeString, ) -> NemoRelayStatus; +/// Native event observability-field sanitizer callback. +pub type NemoRelayNativeEventSanitizeCb = unsafe extern "C" fn( + user_data: *mut c_void, + event_json: *const NemoRelayNativeString, + fields_json: *const NemoRelayNativeString, + out_fields_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus; + /// Native JSON transform callback for tool request/response sanitizers and tool request intercepts. pub type NemoRelayNativeToolJsonCb = unsafe extern "C" fn( user_data: *mut c_void, @@ -495,6 +503,36 @@ pub struct NemoRelayNativeHostApiV1 { cb: NemoRelayNativeWithScopeStackCb, user_data: *mut c_void, ) -> NemoRelayStatus, + /// Registers a mark event sanitizer through the plugin context. + pub plugin_context_register_mark_sanitize_guardrail: unsafe extern "C" fn( + ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeEventSanitizeCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) + -> NemoRelayStatus, + /// Registers a scope-start event sanitizer through the plugin context. + pub plugin_context_register_scope_sanitize_start_guardrail: + unsafe extern "C" fn( + ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeEventSanitizeCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus, + /// Registers a scope-end event sanitizer through the plugin context. + pub plugin_context_register_scope_sanitize_end_guardrail: + unsafe extern "C" fn( + ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeEventSanitizeCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus, } // The host API table is immutable after construction. Function pointers and @@ -1243,6 +1281,97 @@ impl<'a> PluginContext<'a> { finish_typed_registration::(self.host, status, user_data, "subscriber") } + fn register_event_sanitizer( + &mut self, + name: &str, + priority: i32, + callback: F, + register: unsafe extern "C" fn( + *mut NemoRelayNativePluginContext, + *const NemoRelayNativeString, + i32, + NemoRelayNativeEventSanitizeCb, + *mut c_void, + NemoRelayNativeFreeFn, + ) -> NemoRelayStatus, + label: &str, + ) -> Result<()> + where + F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + { + let user_data = typed_callback_user_data(self.host, callback); + let status = self.with_name(name, |_, name| unsafe { + register( + self.raw, + name, + priority, + typed_event_sanitize_trampoline::, + user_data, + Some(drop_typed_callback::), + ) + }); + finish_typed_registration::(self.host, status, user_data, label) + } + + /// Registers a typed mark event sanitizer. + pub fn register_mark_sanitize_guardrail( + &mut self, + name: &str, + priority: i32, + callback: F, + ) -> Result<()> + where + F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + { + self.register_event_sanitizer( + name, + priority, + callback, + self.host.plugin_context_register_mark_sanitize_guardrail, + "mark sanitize guardrail", + ) + } + + /// Registers a typed scope-start event sanitizer. + pub fn register_scope_sanitize_start_guardrail( + &mut self, + name: &str, + priority: i32, + callback: F, + ) -> Result<()> + where + F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + { + self.register_event_sanitizer( + name, + priority, + callback, + self.host + .plugin_context_register_scope_sanitize_start_guardrail, + "scope-start sanitize guardrail", + ) + } + + /// Registers a typed scope-end event sanitizer. + pub fn register_scope_sanitize_end_guardrail( + &mut self, + name: &str, + priority: i32, + callback: F, + ) -> Result<()> + where + F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + { + self.register_event_sanitizer( + name, + priority, + callback, + self.host + .plugin_context_register_scope_sanitize_end_guardrail, + "scope-end sanitize guardrail", + ) + } + /// Registers a typed tool sanitize-request guardrail. pub fn register_tool_sanitize_request_guardrail( &mut self, @@ -1569,6 +1698,69 @@ impl<'a> PluginContext<'a> { }) } + /// Registers a raw mark event sanitizer callback. + /// + /// # Safety + /// `cb`, `user_data`, and `free_fn` must remain valid for every host + /// callback invocation until the host deregisters the callback or calls + /// `free_fn`. `free_fn` must match the allocation behind `user_data`. + pub unsafe fn register_mark_sanitize_guardrail_raw( + &mut self, + name: &str, + priority: i32, + cb: NemoRelayNativeEventSanitizeCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus { + self.with_name(name, |host, name| unsafe { + (host.plugin_context_register_mark_sanitize_guardrail)( + self.raw, name, priority, cb, user_data, free_fn, + ) + }) + } + + /// Registers a raw scope-start event sanitizer callback. + /// + /// # Safety + /// `cb`, `user_data`, and `free_fn` must remain valid for every host + /// callback invocation until the host deregisters the callback or calls + /// `free_fn`. `free_fn` must match the allocation behind `user_data`. + pub unsafe fn register_scope_sanitize_start_guardrail_raw( + &mut self, + name: &str, + priority: i32, + cb: NemoRelayNativeEventSanitizeCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus { + self.with_name(name, |host, name| unsafe { + (host.plugin_context_register_scope_sanitize_start_guardrail)( + self.raw, name, priority, cb, user_data, free_fn, + ) + }) + } + + /// Registers a raw scope-end event sanitizer callback. + /// + /// # Safety + /// `cb`, `user_data`, and `free_fn` must remain valid for every host + /// callback invocation until the host deregisters the callback or calls + /// `free_fn`. `free_fn` must match the allocation behind `user_data`. + pub unsafe fn register_scope_sanitize_end_guardrail_raw( + &mut self, + name: &str, + priority: i32, + cb: NemoRelayNativeEventSanitizeCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus { + self.with_name(name, |host, name| unsafe { + (host.plugin_context_register_scope_sanitize_end_guardrail)( + self.raw, name, priority, cb, user_data, free_fn, + ) + }) + } + /// Registers a raw tool sanitize-request guardrail callback. /// /// # Safety @@ -1902,6 +2094,34 @@ where } } +unsafe extern "C" fn typed_event_sanitize_trampoline( + user_data: *mut c_void, + event_json: *const NemoRelayNativeString, + fields_json: *const NemoRelayNativeString, + out_fields_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus +where + F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, +{ + if user_data.is_null() || out_fields_json.is_null() { + return NemoRelayStatus::NullPointer; + } + unsafe { *out_fields_json = ptr::null_mut() }; + let state = unsafe { &*(user_data as *const TypedCallback) }; + let result = catch_unwind(AssertUnwindSafe(|| { + let event: Event = read_json_value(&state.host, event_json, "event")?; + let fields: EventSanitizeFields = + read_json_value(&state.host, fields_json, "event sanitize fields")?; + let output = (state.callback)(&event, fields); + Ok::<_, NemoRelayStatus>(write_json(&state.host, &output, out_fields_json)) + })); + match result { + Ok(Ok(status)) => status, + Ok(Err(status)) => status, + Err(_) => callback_panic(&state.host, "event sanitize callback"), + } +} + unsafe extern "C" fn typed_tool_sanitize_trampoline( user_data: *mut c_void, name: *const NemoRelayNativeString, diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 194955181..4a1b20041 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -14,10 +14,11 @@ use std::sync::{ use nemo_relay_plugin::{ AnnotatedLlmRequest, CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, - Json, LlmJsonStream, LlmNext, LlmRequest, LlmRequestInterceptOutcome, LlmStream, LlmStreamNext, - NEMO_RELAY_NATIVE_ABI_VERSION, NativePlugin, NemoRelayNativeEventSubscriberCb, - NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeJsonCb, - NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCb, + EventSanitizeFields, Json, LlmJsonStream, LlmNext, LlmRequest, LlmRequestInterceptOutcome, + LlmStream, LlmStreamNext, NEMO_RELAY_NATIVE_ABI_VERSION, NativePlugin, + NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, + NemoRelayNativeHostApiV1, NemoRelayNativeJsonCb, NemoRelayNativeLlmConditionalCb, + NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCb, NemoRelayNativeLlmRequestInterceptCb, NemoRelayNativeLlmStreamExecutionCb, NemoRelayNativeLlmStreamV1, NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, NemoRelayNativeScopeStack, NemoRelayNativeScopeStackBinding, @@ -37,6 +38,22 @@ struct RegisteredSubscriber { free_fn: NemoRelayNativeFreeFn, } +struct RegisteredEventSanitize { + name: String, + priority: i32, + cb: NemoRelayNativeEventSanitizeCb, + user_data: usize, + free_fn: NemoRelayNativeFreeFn, +} + +impl RegisteredEventSanitize { + unsafe fn free(self) { + if let Some(free_fn) = self.free_fn { + unsafe { free_fn(self.user_data as *mut c_void) }; + } + } +} + impl RegisteredSubscriber { unsafe fn free(self) { if let Some(free_fn) = self.free_fn { @@ -209,6 +226,7 @@ macro_rules! impl_captured_registration { impl_captured_registration!( RegisteredSubscriber, + RegisteredEventSanitize, RegisteredToolJson, RegisteredToolConditional, RegisteredToolExecution, @@ -265,6 +283,7 @@ static SCOPE_STACK_FREES: AtomicUsize = AtomicUsize::new(0); static SCOPE_STACK_BINDING_FREES: AtomicUsize = AtomicUsize::new(0); static SCOPE_STACK_BINDING_RESTORES: AtomicUsize = AtomicUsize::new(0); static SUBSCRIBER_REGISTRATION: Mutex> = Mutex::new(None); +static EVENT_SANITIZE_REGISTRATION: Mutex> = Mutex::new(None); static TOOL_JSON_REGISTRATION: Mutex> = Mutex::new(None); static TOOL_CONDITIONAL_REGISTRATION: Mutex> = Mutex::new(None); static TOOL_EXECUTION_REGISTRATION: Mutex> = Mutex::new(None); @@ -297,12 +316,13 @@ fn native_abi_v1_struct_sizes_are_self_describing() { #[cfg(target_pointer_width = "64")] { assert_eq!(align_of::(), 8); - assert_eq!(size_of::(), 272); + assert_eq!(size_of::(), 296); assert_eq!( host_api_offsets(), [ 0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, - 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 264, + 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 264, 272, + 280, 288, ] ); assert_eq!(align_of::(), 8); @@ -316,12 +336,12 @@ fn native_abi_v1_struct_sizes_are_self_describing() { #[cfg(target_pointer_width = "32")] { assert_eq!(align_of::(), 4); - assert_eq!(size_of::(), 136); + assert_eq!(size_of::(), 148); assert_eq!( host_api_offsets(), [ 0, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, - 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, + 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, ] ); assert_eq!(align_of::(), 4); @@ -333,7 +353,7 @@ fn native_abi_v1_struct_sizes_are_self_describing() { } } -fn host_api_offsets() -> [usize; 34] { +fn host_api_offsets() -> [usize; 37] { [ offset_of!(NemoRelayNativeHostApiV1, abi_version), offset_of!(NemoRelayNativeHostApiV1, struct_size), @@ -402,6 +422,18 @@ fn host_api_offsets() -> [usize; 34] { offset_of!(NemoRelayNativeHostApiV1, scope_stack_binding_free), offset_of!(NemoRelayNativeHostApiV1, scope_stack_active), offset_of!(NemoRelayNativeHostApiV1, scope_stack_with_current), + offset_of!( + NemoRelayNativeHostApiV1, + plugin_context_register_mark_sanitize_guardrail + ), + offset_of!( + NemoRelayNativeHostApiV1, + plugin_context_register_scope_sanitize_start_guardrail + ), + offset_of!( + NemoRelayNativeHostApiV1, + plugin_context_register_scope_sanitize_end_guardrail + ), ] } @@ -554,6 +586,15 @@ unsafe extern "C" fn passthrough_tool_json_cb( NemoRelayStatus::Ok } +unsafe extern "C" fn passthrough_event_sanitize_cb( + _user_data: *mut c_void, + _event_json: *const NemoRelayNativeString, + _fields_json: *const NemoRelayNativeString, + _out_fields_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::Ok +} + unsafe extern "C" fn capture_tool_conditional( _ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, @@ -1038,6 +1079,35 @@ unsafe extern "C" fn capture_scope_handle_free(handle: *mut NemoRelayNativeScope SCOPE_HANDLE_FREES.fetch_add(1, Ordering::SeqCst); } } + +unsafe extern "C" fn capture_event_sanitize( + _ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeEventSanitizeCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + let status = *REGISTRATION_STATUS.lock().unwrap(); + if status == NemoRelayStatus::Ok { + let host = test_host(); + let name = match required_host_string(&host, name) { + Ok(name) => name, + Err(status) => return status, + }; + replace_registration( + &EVENT_SANITIZE_REGISTRATION, + RegisteredEventSanitize { + name, + priority, + cb, + user_data: user_data as usize, + free_fn, + }, + ); + } + status +} unsafe extern "C" fn capture_scope_stack_free(stack: *mut NemoRelayNativeScopeStack) { if !stack.is_null() { unsafe { drop(Box::from_raw(stack.cast::())) }; @@ -1092,6 +1162,9 @@ fn test_host() -> NemoRelayNativeHostApiV1 { scope_stack_binding_free: capture_scope_stack_binding_free, scope_stack_active: true_scope_stack_active, scope_stack_with_current: capture_scope_stack_with_current, + plugin_context_register_mark_sanitize_guardrail: capture_event_sanitize, + plugin_context_register_scope_sanitize_start_guardrail: capture_event_sanitize, + plugin_context_register_scope_sanitize_end_guardrail: capture_event_sanitize, } } @@ -1105,6 +1178,7 @@ fn begin_test() -> MutexGuard<'static, ()> { fn reset_state() { clear_registration(&SUBSCRIBER_REGISTRATION); + clear_registration(&EVENT_SANITIZE_REGISTRATION); clear_registration(&TOOL_JSON_REGISTRATION); clear_registration(&TOOL_CONDITIONAL_REGISTRATION); clear_registration(&TOOL_EXECUTION_REGISTRATION); @@ -1296,6 +1370,14 @@ fn take_subscriber_registration() -> RegisteredSubscriber { .expect("subscriber callback should be registered") } +fn take_event_sanitize_registration() -> RegisteredEventSanitize { + EVENT_SANITIZE_REGISTRATION + .lock() + .unwrap() + .take() + .expect("event sanitize callback should be registered") +} + fn take_tool_conditional_registration() -> RegisteredToolConditional { TOOL_CONDITIONAL_REGISTRATION .lock() @@ -2042,6 +2124,81 @@ fn typed_tool_sanitize_guardrails_transform_payloads() { } } +#[test] +fn typed_event_sanitize_guardrails_transform_fields_for_every_surface() { + let _guard = begin_test(); + let host = test_host(); + let mut ctx = test_context(&host); + ctx.register_mark_sanitize_guardrail("mark-sanitize", 4, |event, mut fields| { + assert_eq!(event.name(), "checkpoint"); + fields.data = Some(json!({"clean": true})); + fields.category_profile = None; + fields.metadata = Some(json!({"source": "native"})); + fields + }) + .unwrap(); + + let registration = take_event_sanitize_registration(); + assert_eq!(registration.name, "mark-sanitize"); + assert_eq!(registration.priority, 4); + let event = json_host_string( + &host, + json!({ + "kind": "mark", + "atof_version": "0.1", + "uuid": "00000000-0000-0000-0000-000000000000", + "timestamp": "2026-01-01T00:00:00Z", + "name": "checkpoint", + "data": {"secret": "raw"} + }), + ); + let fields = json_host_string( + &host, + serde_json::to_value(EventSanitizeFields { + data: Some(json!({"secret": "raw"})), + category_profile: Some(CategoryProfile::builder().subtype("raw").build()), + metadata: None, + }) + .unwrap(), + ); + let mut out = ptr::null_mut(); + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + event, + fields, + &mut out, + ) + }; + assert_eq!(status, NemoRelayStatus::Ok); + let output: EventSanitizeFields = + serde_json::from_value(read_json_and_free(&host, out)).unwrap(); + assert_eq!(output.data, Some(json!({"clean": true}))); + assert!(output.category_profile.is_none()); + assert_eq!(output.metadata, Some(json!({"source": "native"}))); + unsafe { + (host.string_free)(event); + (host.string_free)(fields); + registration.free(); + } + + let mut ctx = test_context(&host); + ctx.register_scope_sanitize_start_guardrail("scope-start-sanitize", 5, |_, fields| fields) + .unwrap(); + let registration = take_event_sanitize_registration(); + assert_eq!(registration.name, "scope-start-sanitize"); + assert_eq!(registration.priority, 5); + unsafe { registration.free() }; + + let mut ctx = test_context(&host); + ctx.register_scope_sanitize_end_guardrail("scope-end-sanitize", 6, |_, fields| fields) + .unwrap(); + let registration = take_event_sanitize_registration(); + assert_eq!(registration.name, "scope-end-sanitize"); + assert_eq!(registration.priority, 6); + unsafe { registration.free() }; +} + #[test] fn typed_json_callbacks_report_output_allocation_failures() { let _guard = begin_test(); @@ -2072,6 +2229,84 @@ fn typed_json_callbacks_report_output_allocation_failures() { (host.string_free)(payload); registration.free(); } + + let mut ctx = test_context(&host); + ctx.register_mark_sanitize_guardrail("mark-sanitize", 0, |_, fields| fields) + .unwrap(); + let registration = take_event_sanitize_registration(); + let event = json_host_string( + &host, + json!({ + "kind": "mark", + "atof_version": "0.1", + "uuid": "00000000-0000-0000-0000-000000000000", + "timestamp": "2026-01-01T00:00:00Z", + "name": "checkpoint" + }), + ); + let fields = json_host_string(&host, json!({})); + *STRING_NEW_REMAINING_SUCCESSES.lock().unwrap() = Some(0); + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + event, + fields, + &mut out, + ) + }; + *STRING_NEW_REMAINING_SUCCESSES.lock().unwrap() = None; + assert_eq!(status, NemoRelayStatus::Internal); + assert!(out.is_null()); + unsafe { + (host.string_free)(event); + (host.string_free)(fields); + registration.free(); + } +} + +#[test] +fn typed_event_sanitize_callback_catches_panics() { + let _guard = begin_test(); + let host = test_host(); + let mut ctx = test_context(&host); + ctx.register_mark_sanitize_guardrail("mark-sanitize", 0, |_, _| { + panic!("event sanitizer panic") + }) + .unwrap(); + let registration = take_event_sanitize_registration(); + let event = json_host_string( + &host, + json!({ + "kind": "mark", + "atof_version": "0.1", + "uuid": "00000000-0000-0000-0000-000000000000", + "timestamp": "2026-01-01T00:00:00Z", + "name": "checkpoint" + }), + ); + let fields = json_host_string(&host, json!({})); + let mut out = ptr::null_mut(); + assert_eq!( + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + event, + fields, + &mut out, + ) + }, + NemoRelayStatus::Internal + ); + assert!(out.is_null()); + assert_eq!( + LAST_ERROR.lock().unwrap().as_deref(), + Some("event sanitize callback panicked") + ); + unsafe { + (host.string_free)(event); + (host.string_free)(fields); + registration.free(); + } } #[test] @@ -2389,6 +2624,65 @@ fn typed_callbacks_reject_null_abi_pointers_before_decoding_inputs() { registration.free(); } + let mut ctx = test_context(&host); + ctx.register_mark_sanitize_guardrail("mark-sanitize", 0, |_, fields| fields) + .unwrap(); + let registration = take_event_sanitize_registration(); + let event = json_host_string( + &host, + json!({ + "kind": "mark", + "atof_version": "0.1", + "uuid": "00000000-0000-0000-0000-000000000000", + "timestamp": "2026-01-01T00:00:00Z", + "name": "checkpoint" + }), + ); + let fields = json_host_string(&host, json!({})); + let mut out = ptr::null_mut(); + assert_eq!( + unsafe { (registration.cb)(ptr::null_mut(), event, fields, &mut out) }, + NemoRelayStatus::NullPointer + ); + assert_eq!( + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + ptr::null(), + fields, + &mut out, + ) + }, + NemoRelayStatus::NullPointer + ); + assert_eq!( + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + event, + ptr::null(), + &mut out, + ) + }, + NemoRelayStatus::NullPointer + ); + assert_eq!( + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + event, + fields, + ptr::null_mut(), + ) + }, + NemoRelayStatus::NullPointer + ); + unsafe { + (host.string_free)(event); + (host.string_free)(fields); + registration.free(); + } + let mut ctx = test_context(&host); ctx.register_tool_sanitize_request_guardrail("tool-sanitize", 0, |_name, value| value) .unwrap(); @@ -2733,6 +3027,59 @@ fn typed_callbacks_report_invalid_json_for_each_decoder_family() { registration.free(); } + let mut ctx = test_context(&host); + ctx.register_mark_sanitize_guardrail("mark-sanitize", 0, |_, fields| fields) + .unwrap(); + let registration = take_event_sanitize_registration(); + let invalid_event = host_string(&host, "{not json"); + let fields = json_host_string(&host, json!({})); + let stale_out = host_string(&host, r#"{"stale":true}"#); + let mut out = stale_out; + assert_eq!( + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + invalid_event, + fields, + &mut out, + ) + }, + NemoRelayStatus::InvalidJson + ); + assert!(out.is_null()); + let event = json_host_string( + &host, + json!({ + "kind": "mark", + "atof_version": "0.1", + "uuid": "00000000-0000-0000-0000-000000000000", + "timestamp": "2026-01-01T00:00:00Z", + "name": "checkpoint" + }), + ); + let invalid_fields = host_string(&host, "{not json"); + out = stale_out; + assert_eq!( + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + event, + invalid_fields, + &mut out, + ) + }, + NemoRelayStatus::InvalidJson + ); + assert!(out.is_null()); + unsafe { + (host.string_free)(stale_out); + (host.string_free)(invalid_event); + (host.string_free)(fields); + (host.string_free)(event); + (host.string_free)(invalid_fields); + registration.free(); + } + let mut ctx = test_context(&host); ctx.register_tool_sanitize_request_guardrail("tool-sanitize", 0, |_name, value| value) .unwrap(); @@ -4529,6 +4876,70 @@ fn raw_registration_propagates_name_allocation_status() { assert!(TOOL_JSON_REGISTRATION.lock().unwrap().is_none()); } +#[test] +fn raw_event_sanitize_registrations_cover_every_surface() { + let _guard = begin_test(); + let host = test_host(); + let mut ctx = test_context(&host); + + assert_eq!( + unsafe { + ctx.register_mark_sanitize_guardrail_raw( + "raw-mark", + 1, + passthrough_event_sanitize_cb, + ptr::null_mut(), + None, + ) + }, + NemoRelayStatus::Ok + ); + let registration = take_event_sanitize_registration(); + assert_eq!( + (registration.name.as_str(), registration.priority), + ("raw-mark", 1) + ); + unsafe { registration.free() }; + + assert_eq!( + unsafe { + ctx.register_scope_sanitize_start_guardrail_raw( + "raw-scope-start", + 2, + passthrough_event_sanitize_cb, + ptr::null_mut(), + None, + ) + }, + NemoRelayStatus::Ok + ); + let registration = take_event_sanitize_registration(); + assert_eq!( + (registration.name.as_str(), registration.priority), + ("raw-scope-start", 2) + ); + unsafe { registration.free() }; + + assert_eq!( + unsafe { + ctx.register_scope_sanitize_end_guardrail_raw( + "raw-scope-end", + 3, + passthrough_event_sanitize_cb, + ptr::null_mut(), + None, + ) + }, + NemoRelayStatus::Ok + ); + let registration = take_event_sanitize_registration(); + assert_eq!( + (registration.name.as_str(), registration.priority), + ("raw-scope-end", 3) + ); + unsafe { registration.free() }; +} + #[test] fn typed_registration_name_allocation_failure_drops_callback_state() { let _guard = begin_test(); diff --git a/crates/types/src/api/event.rs b/crates/types/src/api/event.rs index 839419d70..5d8ba788d 100644 --- a/crates/types/src/api/event.rs +++ b/crates/types/src/api/event.rs @@ -364,6 +364,24 @@ pub struct MarkEvent { pub category_profile: Option, } +/// Observability fields that event sanitizers may rewrite. +/// +/// Event identity and lifecycle fields are intentionally excluded so sanitizer +/// callbacks cannot alter correlation, ordering, or category semantics. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, TypedBuilder)] +#[builder(field_defaults(setter(into, strip_option(ignore_invalid, fallback_suffix = "_opt"))))] +pub struct EventSanitizeFields { + /// Application-defined event payload. + #[builder(default)] + pub data: Option, + /// Category-specific typed fields. + #[builder(default)] + pub category_profile: Option, + /// Tracing and correlation metadata. + #[builder(default)] + pub metadata: Option, +} + /// Mark requested by middleware for materialization by a lifecycle owner. /// /// The runtime assigns the parent UUID, event UUID, and timestamp when it @@ -535,6 +553,25 @@ impl Event { self.base().data.as_ref() } + /// Snapshot the observability fields that sanitizers may rewrite. + pub fn sanitize_fields(&self) -> EventSanitizeFields { + EventSanitizeFields { + data: self.base().data.clone(), + category_profile: self.category_profile().cloned(), + metadata: self.base().metadata.clone(), + } + } + + /// Replace the observability fields controlled by event sanitizers. + pub fn apply_sanitize_fields(&mut self, fields: EventSanitizeFields) { + self.base_mut().data = fields.data; + self.base_mut().metadata = fields.metadata; + match self { + Self::Scope(event) => event.category_profile = fields.category_profile, + Self::Mark(event) => event.category_profile = fields.category_profile, + } + } + /// Return the optional data schema. /// /// # Returns @@ -666,6 +703,13 @@ impl Event { Self::Mark(event) => &event.base, } } + + fn base_mut(&mut self) -> &mut BaseEvent { + match self { + Self::Scope(event) => &mut event.base, + Self::Mark(event) => &mut event.base, + } + } } /// Convert handle bitflags into ATOF attributes. diff --git a/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto b/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto index 9dbe9a676..0d17d8b6b 100644 --- a/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto +++ b/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto @@ -46,6 +46,9 @@ enum RegistrationSurface { LLM_REQUEST_INTERCEPT = 23; LLM_EXECUTION_INTERCEPT = 24; LLM_STREAM_EXECUTION_INTERCEPT = 25; + MARK_SANITIZE_GUARDRAIL = 30; + SCOPE_SANITIZE_START_GUARDRAIL = 31; + SCOPE_SANITIZE_END_GUARDRAIL = 32; } enum ScopeType { diff --git a/crates/worker-proto/tests/proto_tests.rs b/crates/worker-proto/tests/proto_tests.rs index 35a478c3e..3b34ad9d1 100644 --- a/crates/worker-proto/tests/proto_tests.rs +++ b/crates/worker-proto/tests/proto_tests.rs @@ -38,6 +38,9 @@ fn registration_surface_values_are_stable() { assert_eq!(RegistrationSurface::LlmRequestIntercept as i32, 23); assert_eq!(RegistrationSurface::LlmExecutionIntercept as i32, 24); assert_eq!(RegistrationSurface::LlmStreamExecutionIntercept as i32, 25); + assert_eq!(RegistrationSurface::MarkSanitizeGuardrail as i32, 30); + assert_eq!(RegistrationSurface::ScopeSanitizeStartGuardrail as i32, 31); + assert_eq!(RegistrationSurface::ScopeSanitizeEndGuardrail as i32, 32); } #[test] diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 56921f544..4a456fc1c 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -19,7 +19,7 @@ //! arbitrary blocking work started by the callback has stopped. use std::cell::RefCell; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::net::{SocketAddr, ToSocketAddrs}; #[cfg(unix)] @@ -34,7 +34,7 @@ use futures_util::{Stream, StreamExt}; #[cfg(unix)] use hyper_util::rt::TokioIo; pub use nemo_relay_types::Json; -pub use nemo_relay_types::api::event::{Event, PendingMarkSpec}; +pub use nemo_relay_types::api::event::{Event, EventSanitizeFields, PendingMarkSpec}; pub use nemo_relay_types::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; pub use nemo_relay_types::api::scope::ScopeType; pub use nemo_relay_types::api::tool::ToolExecutionInterceptOutcome; @@ -118,6 +118,8 @@ pub trait WorkerPlugin: Send + Sync + 'static { } type SubscriberFn = Arc; +type EventSanitizeFn = + Arc EventSanitizeFields + Send + Sync>; type ToolSanitizeFn = Arc Json + Send + Sync>; type ToolConditionalFn = Arc Result> + Send + Sync>; type ToolRequestFn = Arc Result + Send + Sync>; @@ -140,6 +142,9 @@ type LlmStreamExecutionFn = struct WorkerHandlers { registrations: Vec, subscribers: HashMap, + mark_sanitizers: HashMap, + scope_start_sanitizers: HashMap, + scope_end_sanitizers: HashMap, tool_sanitize_requests: HashMap, tool_sanitize_responses: HashMap, tool_conditionals: HashMap, @@ -192,6 +197,76 @@ impl PluginContext { .insert(name.into(), Arc::new(callback)); } + fn register_event_sanitizer( + &mut self, + name: &str, + priority: i32, + surface: RegistrationSurface, + callback: F, + ) where + F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + { + self.push_registration(name, surface, priority, false); + let sanitizers = match surface { + RegistrationSurface::MarkSanitizeGuardrail => &mut self.handlers.mark_sanitizers, + RegistrationSurface::ScopeSanitizeStartGuardrail => { + &mut self.handlers.scope_start_sanitizers + } + RegistrationSurface::ScopeSanitizeEndGuardrail => { + &mut self.handlers.scope_end_sanitizers + } + _ => unreachable!("event sanitizer registration requires an event sanitizer surface"), + }; + sanitizers.insert(name.into(), Arc::new(callback)); + } + + /// Registers a mark event sanitizer. + pub fn register_mark_sanitize_guardrail(&mut self, name: &str, priority: i32, callback: F) + where + F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + { + self.register_event_sanitizer( + name, + priority, + RegistrationSurface::MarkSanitizeGuardrail, + callback, + ); + } + + /// Registers a scope-start event sanitizer. + pub fn register_scope_sanitize_start_guardrail( + &mut self, + name: &str, + priority: i32, + callback: F, + ) where + F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + { + self.register_event_sanitizer( + name, + priority, + RegistrationSurface::ScopeSanitizeStartGuardrail, + callback, + ); + } + + /// Registers a scope-end event sanitizer. + pub fn register_scope_sanitize_end_guardrail( + &mut self, + name: &str, + priority: i32, + callback: F, + ) where + F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + { + self.register_event_sanitizer( + name, + priority, + RegistrationSurface::ScopeSanitizeEndGuardrail, + callback, + ); + } + /// Registers a tool sanitize-request guardrail. pub fn register_tool_sanitize_request_guardrail( &mut self, @@ -968,6 +1043,12 @@ impl PluginWorker for WorkerService { error: Some(sdk_error_to_worker(err)), })); } + if let Err(err) = validate_unique_registrations(&ctx.handlers.registrations) { + return Ok(Response::new(RegisterResponse { + registrations: Vec::new(), + error: Some(sdk_error_to_worker(err)), + })); + } let registrations = ctx.handlers.registrations.clone(); *self .handlers @@ -1230,6 +1311,21 @@ impl PluginWorker for WorkerService { } } +fn validate_unique_registrations(registrations: &[Registration]) -> Result<()> { + let mut seen = HashSet::new(); + for registration in registrations { + if !seen.insert((registration.surface, registration.local_name.as_str())) { + let surface = RegistrationSurface::try_from(registration.surface) + .map_or("UNKNOWN", |surface| surface.as_str_name()); + return Err(WorkerSdkError::InvalidInput(format!( + "duplicate registration '{}' for surface {}", + registration.local_name, surface + ))); + } + } + Ok(()) +} + impl WorkerService { fn authorize(&self, activation_id: &str, auth_token: &str) -> std::result::Result<(), Status> { if activation_id != self.runtime.activation_id { @@ -1334,6 +1430,18 @@ impl WorkerService { with_thread_scope(&scope, || handler(&event)); Ok(empty_response()) } + RegistrationSurface::MarkSanitizeGuardrail + | RegistrationSurface::ScopeSanitizeStartGuardrail + | RegistrationSurface::ScopeSanitizeEndGuardrail => { + let event = event_payload(request.payload)?; + let fields = event.sanitize_fields(); + let handler = self.event_sanitizer(surface, &request.registration_name)?; + let fields = with_thread_scope(&scope, || handler(&event, fields)); + Ok(json_response( + serde_json::to_value(fields) + .expect("event sanitize fields are JSON serializable"), + )) + } RegistrationSurface::ToolSanitizeRequestGuardrail => { let payload = tool_payload(request.payload)?; let handler = self.tool_sanitize_request(&request.registration_name)?; @@ -1443,6 +1551,22 @@ impl WorkerService { }) } + fn event_sanitizer(&self, surface: RegistrationSurface, name: &str) -> Result { + let handlers = self + .handlers + .lock() + .map_err(|err| WorkerSdkError::Callback(format!("handler lock poisoned: {err}")))?; + let sanitizers = match surface { + RegistrationSurface::MarkSanitizeGuardrail => &handlers.mark_sanitizers, + RegistrationSurface::ScopeSanitizeStartGuardrail => &handlers.scope_start_sanitizers, + RegistrationSurface::ScopeSanitizeEndGuardrail => &handlers.scope_end_sanitizers, + _ => unreachable!("event sanitizer lookup requires an event sanitizer surface"), + }; + sanitizers.get(name).cloned().ok_or_else(|| { + WorkerSdkError::InvalidInput(format!("event sanitizer '{name}' not registered")) + }) + } + fn tool_sanitize_request(&self, name: &str) -> Result { self.handlers .lock() @@ -1919,6 +2043,9 @@ fn all_surfaces() -> Vec { RegistrationSurface::LlmRequestIntercept, RegistrationSurface::LlmExecutionIntercept, RegistrationSurface::LlmStreamExecutionIntercept, + RegistrationSurface::MarkSanitizeGuardrail, + RegistrationSurface::ScopeSanitizeStartGuardrail, + RegistrationSurface::ScopeSanitizeEndGuardrail, ] } diff --git a/crates/worker/tests/worker_sdk_tests.rs b/crates/worker/tests/worker_sdk_tests.rs index cbd84de5c..c0443fc86 100644 --- a/crates/worker/tests/worker_sdk_tests.rs +++ b/crates/worker/tests/worker_sdk_tests.rs @@ -198,7 +198,7 @@ async fn worker_service_enforces_auth_and_reports_registrations() { assert_eq!(invalid_register_config.code(), tonic::Code::InvalidArgument); let registrations = register_plugin(&mut client).await; - assert_eq!(registrations.len(), 16); + assert_eq!(registrations.len(), 19); assert_eq!( registrations .iter() @@ -206,6 +206,13 @@ async fn worker_service_enforces_auth_and_reports_registrations() { .count(), 2 ); + assert_eq!( + registrations + .iter() + .filter(|registration| registration.local_name == "event-sanitize") + .count(), + 3 + ); let invoke_err = client .invoke(Request::new(InvokeRequest { @@ -248,6 +255,37 @@ async fn worker_service_enforces_auth_and_reports_registrations() { handle.abort(); } +#[tokio::test(flavor = "multi_thread")] +async fn worker_service_rejects_duplicate_registration_names_on_one_surface() { + let (handle, mut client) = spawn_worker( + Arc::new(DuplicateEventSanitizerPlugin), + "http://127.0.0.1:9".into(), + ) + .await; + + let response = client + .register(Request::new(RegisterRequest { + activation_id: ACTIVATION_ID.into(), + plugin_id: PLUGIN_ID.into(), + auth_token: AUTH_TOKEN.into(), + config: Some(json_env(json!({}))), + })) + .await + .expect("duplicate registration should return protocol data") + .into_inner(); + assert!(response.registrations.is_empty()); + let error = response + .error + .expect("duplicate registration should return an error"); + assert_eq!(error.code, "worker.error"); + assert_eq!( + error.message, + "invalid input: duplicate registration 'duplicate' for surface MARK_SANITIZE_GUARDRAIL" + ); + + handle.abort(); +} + #[tokio::test(flavor = "multi_thread")] async fn worker_service_cancels_unary_and_stream_invocations_by_id() { let timeout = Duration::from_secs(2); @@ -454,6 +492,31 @@ async fn worker_service_invokes_every_registration_surface() { ["subscriber-event"] ); + let mark_fields = invoke_json( + &mut client, + event_invoke_surface("event-sanitize", RegistrationSurface::MarkSanitizeGuardrail), + ) + .await; + assert_eq!(mark_fields["data"]["phase"], "mark"); + let start_fields = invoke_json( + &mut client, + event_invoke_surface( + "event-sanitize", + RegistrationSurface::ScopeSanitizeStartGuardrail, + ), + ) + .await; + assert_eq!(start_fields["metadata"]["phase"], "scope_start"); + let end_fields = invoke_json( + &mut client, + event_invoke_surface( + "event-sanitize", + RegistrationSurface::ScopeSanitizeEndGuardrail, + ), + ) + .await; + assert_eq!(end_fields["metadata"]["phase"], "scope_end"); + assert_json_field( invoke_json( &mut client, @@ -975,6 +1038,27 @@ async fn worker_service_reports_missing_handlers_and_malformed_payloads() { for (request, expected) in [ (event_invoke("missing-subscriber"), "subscriber"), + ( + event_invoke_surface( + "missing-mark-sanitize", + RegistrationSurface::MarkSanitizeGuardrail, + ), + "event sanitizer", + ), + ( + event_invoke_surface( + "missing-scope-start-sanitize", + RegistrationSurface::ScopeSanitizeStartGuardrail, + ), + "event sanitizer", + ), + ( + event_invoke_surface( + "missing-scope-end-sanitize", + RegistrationSurface::ScopeSanitizeEndGuardrail, + ), + "event sanitizer", + ), ( tool_invoke( "missing-tool-sanitize-request", @@ -1087,6 +1171,17 @@ async fn worker_service_reports_missing_handlers_and_malformed_payloads() { .into_inner(), "expected event payload", ); + assert_worker_error( + client + .invoke(Request::new(InvokeRequest { + payload: None, + ..event_invoke_surface("mark-sanitize", RegistrationSurface::MarkSanitizeGuardrail) + })) + .await + .expect("missing sanitizer event payload returns structured error") + .into_inner(), + "expected event payload", + ); assert_worker_error( client .invoke(Request::new(InvokeRequest { @@ -1371,6 +1466,20 @@ impl WorkerPlugin for MinimalPlugin { } } +struct DuplicateEventSanitizerPlugin; + +impl WorkerPlugin for DuplicateEventSanitizerPlugin { + fn plugin_id(&self) -> &str { + PLUGIN_ID + } + + fn register(&self, ctx: &mut PluginContext, _config: &Json) -> Result<()> { + ctx.register_mark_sanitize_guardrail("duplicate", 0, |_, fields| fields); + ctx.register_mark_sanitize_guardrail("duplicate", 1, |_, fields| fields); + Ok(()) + } +} + #[derive(Default)] struct CancellationPlugin { unary_started: Arc, @@ -1503,6 +1612,18 @@ impl WorkerPlugin for SurfacePlugin { .expect("events lock") .push(event.name().into()); }); + ctx.register_mark_sanitize_guardrail("event-sanitize", 1, |event, mut fields| { + fields.data = Some(json!({"name": event.name(), "phase": "mark"})); + fields + }); + ctx.register_scope_sanitize_start_guardrail("event-sanitize", 1, |_, mut fields| { + fields.metadata = Some(json!({"phase": "scope_start"})); + fields + }); + ctx.register_scope_sanitize_end_guardrail("event-sanitize", 1, |_, mut fields| { + fields.metadata = Some(json!({"phase": "scope_end"})); + fields + }); ctx.register_tool_sanitize_request_guardrail("tool-sanitize", 1, |_, value| { set_json_field(value, "phase", "tool_sanitize_request") }); @@ -2104,8 +2225,15 @@ fn invalid_json_env(schema: &str) -> JsonEnvelope { } fn event_invoke(registration_name: &str) -> InvokeRequest { + event_invoke_surface(registration_name, RegistrationSurface::Subscriber) +} + +fn event_invoke_surface(registration_name: &str, surface: RegistrationSurface) -> InvokeRequest { let event = Event::Mark(MarkEvent::new( - BaseEvent::builder().name("subscriber-event").build(), + BaseEvent::builder() + .name("subscriber-event") + .data(json!({"raw": true})) + .build(), None, None, )); @@ -2113,7 +2241,7 @@ fn event_invoke(registration_name: &str) -> InvokeRequest { activation_id: ACTIVATION_ID.into(), invocation_id: "invoke-1".into(), registration_name: registration_name.into(), - surface: RegistrationSurface::Subscriber as i32, + surface: surface as i32, continuation_id: "next-1".into(), scope: Some(scope_context()), auth_token: AUTH_TOKEN.into(),