diff --git a/crates/ffi/nemo_relay.h b/crates/ffi/nemo_relay.h index 24dbffa15..6355d579c 100644 --- a/crates/ffi/nemo_relay.h +++ b/crates/ffi/nemo_relay.h @@ -206,10 +206,12 @@ typedef struct Option_NemoRelayFinalizerCb Option_NemoRelayFinalizerCb; typedef struct Option_NemoRelayPluginValidateCb Option_NemoRelayPluginValidateCb; /** - * Callback for LLM execution (default callable). Receives a native JSON C string, - * returns the response as a JSON C string. + * Callback for mark and scope event sanitizers. + * The returned JSON string transfers to Relay and is freed exactly once. */ -typedef char *(*NemoRelayLlmExecCb)(void *user_data, const char *native_json); +typedef char *(*NemoRelayEventSanitizeCb)(void *user_data, + const struct FfiEvent *event, + const char *fields_json); /** * Optional destructor for user data passed to callbacks. @@ -217,6 +219,12 @@ typedef char *(*NemoRelayLlmExecCb)(void *user_data, const char *native_json); */ typedef void (*NemoRelayFreeFn)(void *user_data); +/** + * Callback for LLM execution (default callable). Receives a native JSON C string, + * returns the response as a JSON C string. + */ +typedef char *(*NemoRelayLlmExecCb)(void *user_data, const char *native_json); + /** * Nullable version of [`NemoRelayCodecDecodeCb`] for use as an optional * parameter in FFI execute functions. Pass null to indicate no codec. @@ -573,6 +581,120 @@ NemoRelayStatus nemo_relay_adaptive_build_cache_telemetry_event(const char *opti */ NemoRelayStatus nemo_relay_adaptive_set_latency_sensitivity(uint32_t value); +/** + * Register a global mark event sanitizer. + * # Safety + * Pointers must be valid for the documented call lifetime. + */ +NemoRelayStatus nemo_relay_register_mark_sanitize_guardrail(const char *name, + int32_t priority, + NemoRelayEventSanitizeCb cb, + void *user_data, + NemoRelayFreeFn free_fn); + +/** + * Deregister a global mark event sanitizer. + * # Safety + * `name` must be a valid C string. + */ +NemoRelayStatus nemo_relay_deregister_mark_sanitize_guardrail(const char *name); + +/** + * Register a global scope-start event sanitizer. + * # Safety + * Pointers must be valid for the documented call lifetime. + */ +NemoRelayStatus nemo_relay_register_scope_sanitize_start_guardrail(const char *name, + int32_t priority, + NemoRelayEventSanitizeCb cb, + void *user_data, + NemoRelayFreeFn free_fn); + +/** + * Deregister a global scope-start event sanitizer. + * # Safety + * `name` must be a valid C string. + */ +NemoRelayStatus nemo_relay_deregister_scope_sanitize_start_guardrail(const char *name); + +/** + * Register a global scope-end event sanitizer. + * # Safety + * Pointers must be valid for the documented call lifetime. + */ +NemoRelayStatus nemo_relay_register_scope_sanitize_end_guardrail(const char *name, + int32_t priority, + NemoRelayEventSanitizeCb cb, + void *user_data, + NemoRelayFreeFn free_fn); + +/** + * Deregister a global scope-end event sanitizer. + * # Safety + * `name` must be a valid C string. + */ +NemoRelayStatus nemo_relay_deregister_scope_sanitize_end_guardrail(const char *name); + +/** + * Register a scope-local mark event sanitizer. + * # Safety + * Pointers must be valid for the documented call lifetime. + */ +NemoRelayStatus nemo_relay_scope_register_mark_sanitize_guardrail(const char *scope_uuid, + const char *name, + int32_t priority, + NemoRelayEventSanitizeCb cb, + void *user_data, + NemoRelayFreeFn free_fn); + +/** + * Deregister a scope-local mark event sanitizer. + * # Safety + * String pointers must be valid C strings. + */ +NemoRelayStatus nemo_relay_scope_deregister_mark_sanitize_guardrail(const char *scope_uuid, + const char *name); + +/** + * Register a scope-local scope-start event sanitizer. + * # Safety + * Pointers must be valid for the documented call lifetime. + */ +NemoRelayStatus nemo_relay_scope_register_scope_sanitize_start_guardrail(const char *scope_uuid, + const char *name, + int32_t priority, + NemoRelayEventSanitizeCb cb, + void *user_data, + NemoRelayFreeFn free_fn); + +/** + * Deregister a scope-local scope-start event sanitizer. + * # Safety + * String pointers must be valid C strings. + */ +NemoRelayStatus nemo_relay_scope_deregister_scope_sanitize_start_guardrail(const char *scope_uuid, + const char *name); + +/** + * Register a scope-local scope-end event sanitizer. + * # Safety + * Pointers must be valid for the documented call lifetime. + */ +NemoRelayStatus nemo_relay_scope_register_scope_sanitize_end_guardrail(const char *scope_uuid, + const char *name, + int32_t priority, + NemoRelayEventSanitizeCb cb, + void *user_data, + NemoRelayFreeFn free_fn); + +/** + * Deregister a scope-local scope-end event sanitizer. + * # Safety + * String pointers must be valid C strings. + */ +NemoRelayStatus nemo_relay_scope_deregister_scope_sanitize_end_guardrail(const char *scope_uuid, + const char *name); + /** * Begin a manual LLM call lifecycle span. * @@ -1356,6 +1478,42 @@ NemoRelayStatus nemo_relay_plugin_context_register_subscriber(struct FfiPluginCo void *user_data, NemoRelayFreeFn free_fn); +/** + * Register a mark event sanitizer into a plugin context. + * # Safety + * Pointers must remain valid for the documented call lifetime. + */ +NemoRelayStatus nemo_relay_plugin_context_register_mark_sanitize_guardrail(struct FfiPluginContext *ctx, + const char *name, + int32_t priority, + NemoRelayEventSanitizeCb cb, + void *user_data, + NemoRelayFreeFn free_fn); + +/** + * Register a scope-start event sanitizer into a plugin context. + * # Safety + * Pointers must remain valid for the documented call lifetime. + */ +NemoRelayStatus nemo_relay_plugin_context_register_scope_sanitize_start_guardrail(struct FfiPluginContext *ctx, + const char *name, + int32_t priority, + NemoRelayEventSanitizeCb cb, + void *user_data, + NemoRelayFreeFn free_fn); + +/** + * Register a scope-end event sanitizer into a plugin context. + * # Safety + * Pointers must remain valid for the documented call lifetime. + */ +NemoRelayStatus nemo_relay_plugin_context_register_scope_sanitize_end_guardrail(struct FfiPluginContext *ctx, + const char *name, + int32_t priority, + NemoRelayEventSanitizeCb cb, + void *user_data, + NemoRelayFreeFn free_fn); + /** * Register a tool sanitize-request guardrail into the plugin registration context. * diff --git a/crates/ffi/src/api/event_registry.rs b/crates/ffi/src/api/event_registry.rs new file mode 100644 index 000000000..13af2b8d8 --- /dev/null +++ b/crates/ffi/src/api/event_registry.rs @@ -0,0 +1,301 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::{ + NemoRelayEventSanitizeCb, NemoRelayFreeFn, NemoRelayStatus, c_char, c_str_to_string, + clear_last_error, core_registry_api, set_last_error, status_from_error, wrap_event_sanitize_fn, +}; + +#[derive(Clone, Copy)] +pub(crate) enum Surface { + Mark, + Start, + End, +} + +unsafe fn register_global( + name: *const c_char, + priority: i32, + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, + surface: Surface, +) -> NemoRelayStatus { + clear_last_error(); + let callback = wrap_event_sanitize_fn(cb, user_data, free_fn); + let name = match c_str_to_string(name) { + Ok(value) => value, + Err(status) => return status, + }; + let result = match surface { + Surface::Mark => { + core_registry_api::register_mark_sanitize_guardrail(&name, priority, callback) + } + Surface::Start => { + core_registry_api::register_scope_sanitize_start_guardrail(&name, priority, callback) + } + Surface::End => { + core_registry_api::register_scope_sanitize_end_guardrail(&name, priority, callback) + } + }; + result + .map(|()| NemoRelayStatus::Ok) + .unwrap_or_else(|error| status_from_error(&error)) +} + +unsafe fn deregister_global(name: *const c_char, surface: Surface) -> NemoRelayStatus { + clear_last_error(); + let name = match c_str_to_string(name) { + Ok(value) => value, + Err(status) => return status, + }; + let result = match surface { + Surface::Mark => core_registry_api::deregister_mark_sanitize_guardrail(&name), + Surface::Start => core_registry_api::deregister_scope_sanitize_start_guardrail(&name), + Surface::End => core_registry_api::deregister_scope_sanitize_end_guardrail(&name), + }; + result + .map(|_| NemoRelayStatus::Ok) + .unwrap_or_else(|error| status_from_error(&error)) +} + +fn parse_scope_uuid(value: *const c_char) -> Result { + let value = c_str_to_string(value)?; + uuid::Uuid::parse_str(&value).map_err(|error| { + set_last_error(&format!("invalid scope UUID: {error}")); + NemoRelayStatus::InvalidArg + }) +} + +unsafe fn register_scope( + scope_uuid: *const c_char, + name: *const c_char, + priority: i32, + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, + surface: Surface, +) -> NemoRelayStatus { + clear_last_error(); + let callback = wrap_event_sanitize_fn(cb, user_data, free_fn); + let uuid = match parse_scope_uuid(scope_uuid) { + Ok(value) => value, + Err(status) => return status, + }; + let name = match c_str_to_string(name) { + Ok(value) => value, + Err(status) => return status, + }; + let result = match surface { + Surface::Mark => core_registry_api::scope_register_mark_sanitize_guardrail( + &uuid, &name, priority, callback, + ), + Surface::Start => core_registry_api::scope_register_scope_sanitize_start_guardrail( + &uuid, &name, priority, callback, + ), + Surface::End => core_registry_api::scope_register_scope_sanitize_end_guardrail( + &uuid, &name, priority, callback, + ), + }; + result + .map(|()| NemoRelayStatus::Ok) + .unwrap_or_else(|error| status_from_error(&error)) +} + +unsafe fn deregister_scope( + scope_uuid: *const c_char, + name: *const c_char, + surface: Surface, +) -> NemoRelayStatus { + clear_last_error(); + let uuid = match parse_scope_uuid(scope_uuid) { + Ok(value) => value, + Err(status) => return status, + }; + let name = match c_str_to_string(name) { + Ok(value) => value, + Err(status) => return status, + }; + let result = match surface { + Surface::Mark => core_registry_api::scope_deregister_mark_sanitize_guardrail(&uuid, &name), + Surface::Start => { + core_registry_api::scope_deregister_scope_sanitize_start_guardrail(&uuid, &name) + } + Surface::End => { + core_registry_api::scope_deregister_scope_sanitize_end_guardrail(&uuid, &name) + } + }; + result + .map(|_| NemoRelayStatus::Ok) + .unwrap_or_else(|error| status_from_error(&error)) +} + +/// Register a global mark event sanitizer. +/// # Safety +/// Pointers must be valid for the documented call lifetime. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_register_mark_sanitize_guardrail( + name: *const c_char, + priority: i32, + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> NemoRelayStatus { + unsafe { register_global(name, priority, cb, user_data, free_fn, Surface::Mark) } +} +/// Deregister a global mark event sanitizer. +/// # Safety +/// `name` must be a valid C string. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_deregister_mark_sanitize_guardrail( + name: *const c_char, +) -> NemoRelayStatus { + unsafe { deregister_global(name, Surface::Mark) } +} +/// Register a global scope-start event sanitizer. +/// # Safety +/// Pointers must be valid for the documented call lifetime. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_register_scope_sanitize_start_guardrail( + name: *const c_char, + priority: i32, + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> NemoRelayStatus { + unsafe { register_global(name, priority, cb, user_data, free_fn, Surface::Start) } +} +/// Deregister a global scope-start event sanitizer. +/// # Safety +/// `name` must be a valid C string. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_deregister_scope_sanitize_start_guardrail( + name: *const c_char, +) -> NemoRelayStatus { + unsafe { deregister_global(name, Surface::Start) } +} +/// Register a global scope-end event sanitizer. +/// # Safety +/// Pointers must be valid for the documented call lifetime. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_register_scope_sanitize_end_guardrail( + name: *const c_char, + priority: i32, + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> NemoRelayStatus { + unsafe { register_global(name, priority, cb, user_data, free_fn, Surface::End) } +} +/// Deregister a global scope-end event sanitizer. +/// # Safety +/// `name` must be a valid C string. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_deregister_scope_sanitize_end_guardrail( + name: *const c_char, +) -> NemoRelayStatus { + unsafe { deregister_global(name, Surface::End) } +} + +/// Register a scope-local mark event sanitizer. +/// # Safety +/// Pointers must be valid for the documented call lifetime. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_scope_register_mark_sanitize_guardrail( + scope_uuid: *const c_char, + name: *const c_char, + priority: i32, + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> NemoRelayStatus { + unsafe { + register_scope( + scope_uuid, + name, + priority, + cb, + user_data, + free_fn, + Surface::Mark, + ) + } +} +/// Deregister a scope-local mark event sanitizer. +/// # Safety +/// String pointers must be valid C strings. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_scope_deregister_mark_sanitize_guardrail( + scope_uuid: *const c_char, + name: *const c_char, +) -> NemoRelayStatus { + unsafe { deregister_scope(scope_uuid, name, Surface::Mark) } +} +/// Register a scope-local scope-start event sanitizer. +/// # Safety +/// Pointers must be valid for the documented call lifetime. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_scope_register_scope_sanitize_start_guardrail( + scope_uuid: *const c_char, + name: *const c_char, + priority: i32, + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> NemoRelayStatus { + unsafe { + register_scope( + scope_uuid, + name, + priority, + cb, + user_data, + free_fn, + Surface::Start, + ) + } +} +/// Deregister a scope-local scope-start event sanitizer. +/// # Safety +/// String pointers must be valid C strings. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_scope_deregister_scope_sanitize_start_guardrail( + scope_uuid: *const c_char, + name: *const c_char, +) -> NemoRelayStatus { + unsafe { deregister_scope(scope_uuid, name, Surface::Start) } +} +/// Register a scope-local scope-end event sanitizer. +/// # Safety +/// Pointers must be valid for the documented call lifetime. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_scope_register_scope_sanitize_end_guardrail( + scope_uuid: *const c_char, + name: *const c_char, + priority: i32, + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> NemoRelayStatus { + unsafe { + register_scope( + scope_uuid, + name, + priority, + cb, + user_data, + free_fn, + Surface::End, + ) + } +} +/// Deregister a scope-local scope-end event sanitizer. +/// # Safety +/// String pointers must be valid C strings. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_scope_deregister_scope_sanitize_end_guardrail( + scope_uuid: *const c_char, + name: *const c_char, +) -> NemoRelayStatus { + unsafe { deregister_scope(scope_uuid, name, Surface::End) } +} diff --git a/crates/ffi/src/api/mod.rs b/crates/ffi/src/api/mod.rs index 2efe19a67..1d62d9e69 100644 --- a/crates/ffi/src/api/mod.rs +++ b/crates/ffi/src/api/mod.rs @@ -14,17 +14,17 @@ use std::sync::{Arc, OnceLock}; use std::time::Duration; use crate::callable::{ - NemoRelayCodecDecodeFn, NemoRelayCodecEncodeFn, NemoRelayCollectorCb, + NemoRelayCodecDecodeFn, NemoRelayCodecEncodeFn, NemoRelayCollectorCb, NemoRelayEventSanitizeCb, NemoRelayEventSubscriberCb, NemoRelayFinalizerCb, NemoRelayFreeFn, NemoRelayJsonCb, NemoRelayLlmConditionalCb, NemoRelayLlmExecCb, NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestCb, NemoRelayLlmRequestInterceptCb, NemoRelayPluginRegisterCb, NemoRelayPluginValidateCb, NemoRelayToolConditionalCb, NemoRelayToolExecCb, NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, wrap_codec_fn, wrap_collector_fn, - wrap_event_subscriber, wrap_finalizer_fn, wrap_llm_conditional_fn, wrap_llm_exec_fn, - wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, wrap_llm_response_fn, - wrap_llm_sanitize_request_fn, wrap_llm_stream_exec_fn, wrap_llm_stream_exec_intercept_fn, - wrap_tool_conditional_fn, wrap_tool_exec_fn, wrap_tool_exec_intercept_fn, - wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, + wrap_event_sanitize_fn, wrap_event_subscriber, wrap_finalizer_fn, wrap_llm_conditional_fn, + wrap_llm_exec_fn, wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, + wrap_llm_response_fn, wrap_llm_sanitize_request_fn, wrap_llm_stream_exec_fn, + wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, wrap_tool_exec_fn, + wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, }; use crate::convert::{ c_str_to_json, c_str_to_opt_json, c_str_to_string, json_to_c_string, nemo_relay_string_free, @@ -66,6 +66,7 @@ use nemo_relay_adaptive::plugin_component::register_adaptive_component; use tokio::runtime::Runtime; mod adaptive; +mod event_registry; mod llm; mod llm_registry; mod observability; @@ -77,6 +78,7 @@ mod tool_lifecycle; mod tool_registry; pub use adaptive::*; +pub use event_registry::*; pub use llm::*; pub use llm_registry::*; pub use observability::*; diff --git a/crates/ffi/src/api/plugin.rs b/crates/ffi/src/api/plugin.rs index f8dd4752f..51021e09e 100644 --- a/crates/ffi/src/api/plugin.rs +++ b/crates/ffi/src/api/plugin.rs @@ -3,20 +3,21 @@ use super::{ Arc, CStr, ConfigDiagnostic, DiagnosticLevel, FfiPluginContext, Future, - NemoRelayEventSubscriberCb, NemoRelayFreeFn, NemoRelayJsonCb, NemoRelayLlmConditionalCb, - NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestCb, NemoRelayLlmRequestInterceptCb, - NemoRelayPluginRegisterCb, NemoRelayPluginValidateCb, NemoRelayStatus, - NemoRelayToolConditionalCb, NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, Pin, Plugin, - PluginConfig, PluginError, PluginRegistrationContext, active_plugin_report, c_char, - c_str_to_json, c_str_to_string, clear_last_error, clear_plugin_configuration, - deregister_plugin, initialize_plugins, json_to_c_string, last_error_message, list_plugin_kinds, - nemo_relay_string_free, register_adaptive_component, register_plugin, set_last_error, - status_from_plugin_error, tokio_runtime, validate_plugin_config, wrap_event_subscriber, - wrap_llm_conditional_fn, wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, - wrap_llm_response_fn, wrap_llm_sanitize_request_fn, wrap_llm_stream_exec_intercept_fn, - wrap_tool_conditional_fn, wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, - wrap_tool_sanitize_fn, + NemoRelayEventSanitizeCb, NemoRelayEventSubscriberCb, NemoRelayFreeFn, NemoRelayJsonCb, + NemoRelayLlmConditionalCb, NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestCb, + NemoRelayLlmRequestInterceptCb, NemoRelayPluginRegisterCb, NemoRelayPluginValidateCb, + NemoRelayStatus, NemoRelayToolConditionalCb, NemoRelayToolExecInterceptCb, + NemoRelayToolSanitizeCb, Pin, Plugin, PluginConfig, PluginError, PluginRegistrationContext, + active_plugin_report, c_char, c_str_to_json, c_str_to_string, clear_last_error, + clear_plugin_configuration, deregister_plugin, initialize_plugins, json_to_c_string, + last_error_message, list_plugin_kinds, nemo_relay_string_free, register_adaptive_component, + register_plugin, set_last_error, status_from_plugin_error, tokio_runtime, + validate_plugin_config, wrap_event_sanitize_fn, wrap_event_subscriber, wrap_llm_conditional_fn, + wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, wrap_llm_response_fn, + wrap_llm_sanitize_request_fn, wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, + wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, }; +use crate::api::event_registry::Surface; use nemo_relay_pii_redaction::component::register_pii_redaction_component; struct FfiHostedPluginUserData { @@ -366,6 +367,89 @@ pub unsafe extern "C" fn nemo_relay_plugin_context_register_subscriber( } } +unsafe fn plugin_register_event_sanitizer( + ctx: *mut FfiPluginContext, + name: *const c_char, + priority: i32, + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, + surface: Surface, +) -> NemoRelayStatus { + clear_last_error(); + let callback = wrap_event_sanitize_fn(cb, user_data, free_fn); + if ctx.is_null() { + set_last_error("plugin context is null"); + return NemoRelayStatus::NullPointer; + } + let name = match c_str_to_string(name) { + Ok(value) => value, + Err(status) => return status, + }; + let context = unsafe { &mut *((*ctx).0) }; + let result = match surface { + Surface::Mark => context.register_mark_sanitize_guardrail(&name, priority, callback), + Surface::Start => { + context.register_scope_sanitize_start_guardrail(&name, priority, callback) + } + Surface::End => context.register_scope_sanitize_end_guardrail(&name, priority, callback), + }; + result + .map(|()| NemoRelayStatus::Ok) + .unwrap_or_else(|error| status_from_plugin_error(&error)) +} + +/// Register a mark event sanitizer into a plugin context. +/// # Safety +/// Pointers must remain valid for the documented call lifetime. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_plugin_context_register_mark_sanitize_guardrail( + ctx: *mut FfiPluginContext, + name: *const c_char, + priority: i32, + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> NemoRelayStatus { + unsafe { + plugin_register_event_sanitizer(ctx, name, priority, cb, user_data, free_fn, Surface::Mark) + } +} + +/// Register a scope-start event sanitizer into a plugin context. +/// # Safety +/// Pointers must remain valid for the documented call lifetime. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_plugin_context_register_scope_sanitize_start_guardrail( + ctx: *mut FfiPluginContext, + name: *const c_char, + priority: i32, + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> NemoRelayStatus { + unsafe { + plugin_register_event_sanitizer(ctx, name, priority, cb, user_data, free_fn, Surface::Start) + } +} + +/// Register a scope-end event sanitizer into a plugin context. +/// # Safety +/// Pointers must remain valid for the documented call lifetime. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_plugin_context_register_scope_sanitize_end_guardrail( + ctx: *mut FfiPluginContext, + name: *const c_char, + priority: i32, + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> NemoRelayStatus { + unsafe { + plugin_register_event_sanitizer(ctx, name, priority, cb, user_data, free_fn, Surface::End) + } +} + /// Register a tool sanitize-request guardrail into the plugin registration context. /// /// # Safety diff --git a/crates/ffi/src/callable.rs b/crates/ffi/src/callable.rs index 26cb4000a..955e8fed7 100644 --- a/crates/ffi/src/callable.rs +++ b/crates/ffi/src/callable.rs @@ -23,14 +23,14 @@ use std::sync::Arc; use libc::c_char; use nemo_relay::api::runtime::{ - EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmRequestInterceptFn, - LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, - ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, + LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, + ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use serde_json::Value as Json; use tokio_stream::{Stream, StreamExt}; -use nemo_relay::api::event::Event; +use nemo_relay::api::event::{Event, EventSanitizeFields}; use nemo_relay::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; use nemo_relay::api::tool::ToolExecutionInterceptOutcome; use nemo_relay::codec::request::AnnotatedLlmRequest as AnnotatedLLMRequest; @@ -140,6 +140,14 @@ pub type NemoRelayLlmExecInterceptCb = unsafe extern "C" fn( pub type NemoRelayEventSubscriberCb = unsafe extern "C" fn(user_data: *mut libc::c_void, event: *const FfiEvent); +/// Callback for mark and scope event sanitizers. +/// The returned JSON string transfers to Relay and is freed exactly once. +pub type NemoRelayEventSanitizeCb = unsafe extern "C" fn( + user_data: *mut libc::c_void, + event: *const FfiEvent, + fields_json: *const c_char, +) -> *mut c_char; + /// Callback for Codec decode: translates an opaque `FfiLLMRequest` into /// an `AnnotatedLLMRequest` JSON string. Returns a heap-allocated C string /// on success, or null on error (after setting the last error message). @@ -826,6 +834,24 @@ pub fn wrap_event_subscriber( }) } +/// Wrap a C event sanitizer callback into a Rust closure. +pub fn wrap_event_sanitize_fn( + cb: NemoRelayEventSanitizeCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> EventSanitizeFn { + let ud = make_user_data(user_data, free_fn); + Arc::new(move |event: &Event, fields: EventSanitizeFields| { + let ffi_event = FfiEvent(event.clone()); + let fields_json = json_to_c_string(&serde_json::to_value(&fields).unwrap_or(Json::Null)); + let result_ptr = unsafe { cb(ud.ptr, &ffi_event, fields_json) }; + unsafe { nemo_relay_string_free_internal(fields_json) }; + let result = serde_json::from_value(ptr_to_json(result_ptr)).unwrap_or(fields); + unsafe { nemo_relay_string_free_internal(result_ptr) }; + result + }) +} + // --------------------------------------------------------------------------- // Codec wrapper: C callbacks -> Arc // --------------------------------------------------------------------------- diff --git a/crates/ffi/tests/unit/api/registry_tests.rs b/crates/ffi/tests/unit/api/registry_tests.rs index c0809d2c9..8c9a29471 100644 --- a/crates/ffi/tests/unit/api/registry_tests.rs +++ b/crates/ffi/tests/unit/api/registry_tests.rs @@ -4,6 +4,429 @@ //! Unit tests for registry in the NeMo Relay FFI crate. use super::*; +use nemo_relay::plugin::rollback_registrations; + +unsafe extern "C" fn event_sanitize_cb( + _user_data: *mut libc::c_void, + event: *const FfiEvent, + fields_json: *const c_char, +) -> *mut c_char { + let name = unsafe { take_string(nemo_relay_event_name(event)) }.unwrap_or_default(); + let mut fields: Json = serde_json::from_str( + unsafe { CStr::from_ptr(fields_json) } + .to_str() + .unwrap_or("null"), + ) + .unwrap(); + fields["data"] = json!({"sanitized_by": name}); + fields["category_profile"] = json!({"subtype": "ffi.sanitized"}); + fields["metadata"] = Json::Null; + CString::new(fields.to_string()).unwrap().into_raw() +} + +unsafe extern "C" fn invalid_event_sanitize_cb( + _user_data: *mut libc::c_void, + _event: *const FfiEvent, + _fields_json: *const c_char, +) -> *mut c_char { + CString::new("not-json").unwrap().into_raw() +} + +#[test] +fn test_ffi_event_sanitizer_registries_and_error_paths() { + let _lock = TEST_MUTEX.lock().unwrap_or_else(|e| e.into_inner()); + reset_globals(); + + unsafe { + let stack = fresh_scope_stack(); + let subscriber_name = cstring(&unique_name("ffi_event_sanitize_subscriber")); + assert_eq!( + nemo_relay_register_subscriber( + subscriber_name.as_ptr(), + subscriber_cb, + ptr::null_mut(), + None, + ), + NemoRelayStatus::Ok + ); + + let mark_guard = cstring(&unique_name("ffi_mark_sanitize")); + let start_guard = cstring(&unique_name("ffi_scope_start_sanitize")); + let end_guard = cstring(&unique_name("ffi_scope_end_sanitize")); + for status in [ + nemo_relay_register_mark_sanitize_guardrail( + mark_guard.as_ptr(), + 1, + event_sanitize_cb, + Box::into_raw(Box::new(1usize)).cast(), + Some(plugin_free), + ), + nemo_relay_register_scope_sanitize_start_guardrail( + start_guard.as_ptr(), + 1, + event_sanitize_cb, + Box::into_raw(Box::new(2usize)).cast(), + Some(plugin_free), + ), + nemo_relay_register_scope_sanitize_end_guardrail( + end_guard.as_ptr(), + 1, + event_sanitize_cb, + Box::into_raw(Box::new(3usize)).cast(), + Some(plugin_free), + ), + ] { + assert_eq!(status, NemoRelayStatus::Ok); + } + + let scope_name = cstring("ffi-global-scope"); + let original = cstring(r#"{"secret":true}"#); + let mut scope = ptr::null_mut(); + assert_eq!( + nemo_relay_push_scope( + scope_name.as_ptr(), + NemoRelayScopeType::Custom, + ptr::null(), + 0, + original.as_ptr(), + original.as_ptr(), + ptr::null(), + &mut scope, + ), + NemoRelayStatus::Ok + ); + let mark_name = cstring("ffi-global-mark"); + assert_eq!( + nemo_relay_event( + mark_name.as_ptr(), + scope, + original.as_ptr(), + original.as_ptr(), + ), + NemoRelayStatus::Ok + ); + assert_eq!( + nemo_relay_pop_scope(scope, ptr::null()), + NemoRelayStatus::Ok + ); + nemo_relay_scope_handle_free(scope); + assert_eq!(nemo_relay_flush_subscribers(), NemoRelayStatus::Ok); + + let events = lock_unpoisoned(event_log()); + for name in ["ffi-global-scope", "ffi-global-mark"] { + for event in events.iter().filter(|event| event["name"] == name) { + assert_eq!(event["data"], json!({"sanitized_by": name})); + assert_eq!( + event["json"]["category_profile"]["subtype"], + "ffi.sanitized" + ); + assert_eq!(event["metadata"], Json::Null); + } + } + for phase in ["start", "end"] { + assert!(events.iter().any(|event| { + event["name"] == "ffi-global-scope" && event["json"]["scope_category"] == phase + })); + } + drop(events); + + assert_eq!( + nemo_relay_deregister_mark_sanitize_guardrail(mark_guard.as_ptr()), + NemoRelayStatus::Ok + ); + assert_eq!( + nemo_relay_deregister_scope_sanitize_start_guardrail(start_guard.as_ptr()), + NemoRelayStatus::Ok + ); + assert_eq!( + nemo_relay_deregister_scope_sanitize_end_guardrail(end_guard.as_ptr()), + NemoRelayStatus::Ok + ); + assert_eq!(*lock_unpoisoned(plugin_frees()), 3); + + let invalid_guard = cstring(&unique_name("ffi_invalid_event_sanitize")); + assert_eq!( + nemo_relay_register_mark_sanitize_guardrail( + invalid_guard.as_ptr(), + 1, + invalid_event_sanitize_cb, + Box::into_raw(Box::new(4usize)).cast(), + Some(plugin_free), + ), + NemoRelayStatus::Ok + ); + let invalid_mark = cstring("ffi-invalid-callback-mark"); + assert_eq!( + nemo_relay_event( + invalid_mark.as_ptr(), + ptr::null(), + original.as_ptr(), + original.as_ptr(), + ), + NemoRelayStatus::Ok + ); + assert_eq!( + nemo_relay_deregister_mark_sanitize_guardrail(invalid_guard.as_ptr()), + NemoRelayStatus::Ok + ); + assert_eq!(*lock_unpoisoned(plugin_frees()), 4); + + let mut owner = ptr::null_mut(); + let owner_name = cstring("ffi-event-owner"); + assert_eq!( + nemo_relay_push_scope( + owner_name.as_ptr(), + NemoRelayScopeType::Agent, + ptr::null(), + 0, + ptr::null(), + ptr::null(), + ptr::null(), + &mut owner, + ), + NemoRelayStatus::Ok + ); + let owner_uuid = cstring(&take_string(nemo_relay_scope_handle_uuid(owner)).unwrap()); + let local_mark = cstring(&unique_name("ffi_local_mark_sanitize")); + let local_start = cstring(&unique_name("ffi_local_start_sanitize")); + let local_end = cstring(&unique_name("ffi_local_end_sanitize")); + for status in [ + nemo_relay_scope_register_mark_sanitize_guardrail( + owner_uuid.as_ptr(), + local_mark.as_ptr(), + 1, + event_sanitize_cb, + Box::into_raw(Box::new(5usize)).cast(), + Some(plugin_free), + ), + nemo_relay_scope_register_scope_sanitize_start_guardrail( + owner_uuid.as_ptr(), + local_start.as_ptr(), + 1, + event_sanitize_cb, + Box::into_raw(Box::new(6usize)).cast(), + Some(plugin_free), + ), + nemo_relay_scope_register_scope_sanitize_end_guardrail( + owner_uuid.as_ptr(), + local_end.as_ptr(), + 1, + event_sanitize_cb, + Box::into_raw(Box::new(7usize)).cast(), + Some(plugin_free), + ), + ] { + assert_eq!(status, NemoRelayStatus::Ok); + } + let child_name = cstring("ffi-local-child"); + let mut child = ptr::null_mut(); + assert_eq!( + nemo_relay_push_scope( + child_name.as_ptr(), + NemoRelayScopeType::Function, + ptr::null(), + 0, + original.as_ptr(), + original.as_ptr(), + ptr::null(), + &mut child, + ), + NemoRelayStatus::Ok + ); + let local_mark_name = cstring("ffi-local-mark"); + assert_eq!( + nemo_relay_event( + local_mark_name.as_ptr(), + child, + original.as_ptr(), + original.as_ptr(), + ), + NemoRelayStatus::Ok + ); + assert_eq!( + nemo_relay_pop_scope(child, ptr::null()), + NemoRelayStatus::Ok + ); + nemo_relay_scope_handle_free(child); + assert_eq!( + nemo_relay_scope_deregister_mark_sanitize_guardrail( + owner_uuid.as_ptr(), + local_mark.as_ptr(), + ), + NemoRelayStatus::Ok + ); + assert_eq!( + nemo_relay_scope_deregister_scope_sanitize_start_guardrail( + owner_uuid.as_ptr(), + local_start.as_ptr(), + ), + NemoRelayStatus::Ok + ); + assert_eq!( + nemo_relay_scope_deregister_scope_sanitize_end_guardrail( + owner_uuid.as_ptr(), + local_end.as_ptr(), + ), + NemoRelayStatus::Ok + ); + assert_eq!(*lock_unpoisoned(plugin_frees()), 7); + + let invalid_uuid = cstring("not-a-uuid"); + let invalid_name = cstring("invalid-scope-event-sanitizer"); + assert_eq!( + nemo_relay_scope_register_mark_sanitize_guardrail( + invalid_uuid.as_ptr(), + invalid_name.as_ptr(), + 1, + event_sanitize_cb, + ptr::null_mut(), + None, + ), + NemoRelayStatus::InvalidArg + ); + assert_eq!( + nemo_relay_scope_deregister_scope_sanitize_start_guardrail( + invalid_uuid.as_ptr(), + invalid_name.as_ptr(), + ), + NemoRelayStatus::InvalidArg + ); + assert_eq!( + nemo_relay_register_scope_sanitize_end_guardrail( + ptr::null(), + 1, + event_sanitize_cb, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + nemo_relay_deregister_mark_sanitize_guardrail(ptr::null()), + NemoRelayStatus::NullPointer + ); + assert_eq!( + nemo_relay_scope_register_scope_sanitize_end_guardrail( + owner_uuid.as_ptr(), + ptr::null(), + 1, + event_sanitize_cb, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + nemo_relay_scope_deregister_scope_sanitize_end_guardrail( + owner_uuid.as_ptr(), + ptr::null(), + ), + NemoRelayStatus::NullPointer + ); + + assert_eq!( + nemo_relay_pop_scope(owner, ptr::null()), + NemoRelayStatus::Ok + ); + nemo_relay_scope_handle_free(owner); + assert_eq!(nemo_relay_flush_subscribers(), NemoRelayStatus::Ok); + let events = lock_unpoisoned(event_log()); + let invalid_callback_event = events + .iter() + .find(|event| event["name"] == "ffi-invalid-callback-mark") + .expect("invalid callback mark should be delivered"); + assert_eq!(invalid_callback_event["data"], json!({"secret": true})); + assert_eq!(invalid_callback_event["metadata"], json!({"secret": true})); + for name in ["ffi-local-child", "ffi-local-mark"] { + for event in events.iter().filter(|event| event["name"] == name) { + assert_eq!(event["data"], json!({"sanitized_by": name})); + assert_eq!( + event["json"]["category_profile"]["subtype"], + "ffi.sanitized" + ); + assert_eq!(event["metadata"], Json::Null); + } + } + for phase in ["start", "end"] { + assert!(events.iter().any(|event| { + event["name"] == "ffi-local-child" && event["json"]["scope_category"] == phase + })); + } + drop(events); + + let mut context = PluginRegistrationContext::with_namespace(unique_name("ffi_plugin_")); + let mut ffi_context = FfiPluginContext(&mut context); + for (name, status) in [ + ( + cstring("mark"), + nemo_relay_plugin_context_register_mark_sanitize_guardrail( + &mut ffi_context, + cstring("mark").as_ptr(), + 1, + event_sanitize_cb, + Box::into_raw(Box::new(8usize)).cast(), + Some(plugin_free), + ), + ), + ( + cstring("start"), + nemo_relay_plugin_context_register_scope_sanitize_start_guardrail( + &mut ffi_context, + cstring("start").as_ptr(), + 1, + event_sanitize_cb, + Box::into_raw(Box::new(9usize)).cast(), + Some(plugin_free), + ), + ), + ( + cstring("end"), + nemo_relay_plugin_context_register_scope_sanitize_end_guardrail( + &mut ffi_context, + cstring("end").as_ptr(), + 1, + event_sanitize_cb, + Box::into_raw(Box::new(10usize)).cast(), + Some(plugin_free), + ), + ), + ] { + assert!(!name.as_bytes().is_empty()); + assert_eq!(status, NemoRelayStatus::Ok); + } + assert_eq!( + nemo_relay_plugin_context_register_mark_sanitize_guardrail( + ptr::null_mut(), + invalid_name.as_ptr(), + 1, + event_sanitize_cb, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + assert_eq!( + nemo_relay_plugin_context_register_mark_sanitize_guardrail( + &mut ffi_context, + ptr::null(), + 1, + event_sanitize_cb, + ptr::null_mut(), + None, + ), + NemoRelayStatus::NullPointer + ); + let mut registrations = context.into_registrations(); + rollback_registrations(&mut registrations); + assert_eq!(*lock_unpoisoned(plugin_frees()), 10); + + assert_eq!( + nemo_relay_deregister_subscriber(subscriber_name.as_ptr()), + NemoRelayStatus::Ok + ); + nemo_relay_scope_stack_free(stack); + } +} #[test] fn test_ffi_open_telemetry_subscriber_lifecycle_and_errors() { diff --git a/crates/ffi/tests/unit/callable_tests.rs b/crates/ffi/tests/unit/callable_tests.rs index c512d7329..4791b2a46 100644 --- a/crates/ffi/tests/unit/callable_tests.rs +++ b/crates/ffi/tests/unit/callable_tests.rs @@ -6,7 +6,7 @@ use super::*; use std::sync::atomic::{AtomicUsize, Ordering}; -use nemo_relay::api::event::Event; +use nemo_relay::api::event::{Event, EventSanitizeFields}; use nemo_relay::api::llm::{LlmAttributes, LlmHandle}; use serde_json::json; use tokio_stream::StreamExt; @@ -244,6 +244,42 @@ unsafe extern "C" fn subscriber_cb(user_data: *mut libc::c_void, event: *const F } } +unsafe extern "C" fn event_sanitize_cb( + user_data: *mut libc::c_void, + event: *const FfiEvent, + fields_json: *const c_char, +) -> *mut c_char { + let counter = unsafe { &*(user_data as *const Arc) }; + counter.fetch_add(1, Ordering::SeqCst); + assert_eq!(unsafe { (&*event).0.name() }, "ffi-event"); + let mut fields: Json = serde_json::from_str( + unsafe { CStr::from_ptr(fields_json) } + .to_str() + .unwrap_or("null"), + ) + .unwrap(); + fields["data"] = json!({"safe": true}); + fields["category_profile"] = json!({"subtype": "ffi"}); + fields["metadata"] = Json::Null; + CString::new(fields.to_string()).unwrap().into_raw() +} + +unsafe extern "C" fn invalid_event_sanitize_cb( + _user_data: *mut libc::c_void, + _event: *const FfiEvent, + _fields_json: *const c_char, +) -> *mut c_char { + CString::new("invalid-json").unwrap().into_raw() +} + +unsafe extern "C" fn null_event_sanitize_cb( + _user_data: *mut libc::c_void, + _event: *const FfiEvent, + _fields_json: *const c_char, +) -> *mut c_char { + std::ptr::null_mut() +} + fn make_request() -> LlmRequest { LlmRequest { headers: serde_json::Map::new(), @@ -496,6 +532,31 @@ fn test_wrap_llm_exec_stream_and_event_callbacks() { drop(subscriber); assert_eq!(seen.load(Ordering::SeqCst), 2); + let original_fields = EventSanitizeFields::builder() + .data(json!({"secret": true})) + .metadata(json!({"trace": true})) + .build(); + let (user_data, sanitize_calls) = user_data_counter(); + let sanitizer = wrap_event_sanitize_fn(event_sanitize_cb, user_data, Some(free_arc_counter)); + let sanitized = sanitizer(&event, original_fields.clone()); + assert_eq!(sanitized.data, Some(json!({"safe": true}))); + assert_eq!( + sanitized + .category_profile + .as_ref() + .and_then(|profile| profile.subtype.as_deref()), + Some("ffi") + ); + assert_eq!(sanitized.metadata, None); + assert_eq!(sanitize_calls.load(Ordering::SeqCst), 1); + drop(sanitizer); + assert_eq!(sanitize_calls.load(Ordering::SeqCst), 2); + + let invalid = wrap_event_sanitize_fn(invalid_event_sanitize_cb, std::ptr::null_mut(), None); + assert_eq!(invalid(&event, original_fields.clone()), original_fields); + let null = wrap_event_sanitize_fn(null_event_sanitize_cb, std::ptr::null_mut(), None); + assert_eq!(null(&event, original_fields.clone()), original_fields); + let handle = LlmHandle::builder() .name("llm") .attributes(LlmAttributes::STATEFUL) diff --git a/go/nemo_relay/adaptive_plugin_test.go b/go/nemo_relay/adaptive_plugin_test.go index 025ae70f6..75687c099 100644 --- a/go/nemo_relay/adaptive_plugin_test.go +++ b/go/nemo_relay/adaptive_plugin_test.go @@ -6,6 +6,7 @@ package nemo_relay import ( "encoding/json" "fmt" + "sync" "testing" ) @@ -41,6 +42,24 @@ func registerLifecycleGuardrails(ctx *PluginContext) error { }); err != nil { return err } + if err := ctx.RegisterMarkSanitizeGuardrail("mark_sanitize", 7, func(_ Event, fields EventSanitizeFields) EventSanitizeFields { + fields.Data = json.RawMessage(`{"plugin":true}`) + return fields + }); err != nil { + return err + } + if err := ctx.RegisterScopeSanitizeStartGuardrail("scope_start_sanitize", 7, func(_ Event, fields EventSanitizeFields) EventSanitizeFields { + fields.Metadata = json.RawMessage(`{"plugin_phase":"start"}`) + return fields + }); err != nil { + return err + } + if err := ctx.RegisterScopeSanitizeEndGuardrail("scope_end_sanitize", 7, func(_ Event, fields EventSanitizeFields) EventSanitizeFields { + fields.Metadata = json.RawMessage(`{"plugin_phase":"end"}`) + return fields + }); err != nil { + return err + } if err := ctx.RegisterToolSanitizeRequestGuardrail( "tool_sanitize_request", 7, @@ -352,6 +371,56 @@ func TestTopLevelPluginValidationAndLifecycle(t *testing.T) { assertGuardrailError(t, ToolConditionalExecution("blocked-tool", json.RawMessage(`{}`)), "guardrail rejected: blocked tool") assertGuardrailError(t, LlmConditionalExecution(json.RawMessage(`{"headers":{"blocked":true},"content":{"messages":[]}}`)), "guardrail rejected: blocked llm") assertStreamPluginPayload(t, streamPluginKind) + + var mu sync.Mutex + var pluginEvents []Event + if err := RegisterSubscriber("go-plugin-event-sanitize-test", func(event Event) { + mu.Lock() + pluginEvents = append(pluginEvents, event) + mu.Unlock() + }); err != nil { + t.Fatalf("RegisterSubscriber failed: %v", err) + } + defer DeregisterSubscriber("go-plugin-event-sanitize-test") + if err := EmitEvent("plugin-mark", WithEventData(json.RawMessage(`{"raw":true}`))); err != nil { + t.Fatalf("EmitEvent failed: %v", err) + } + handle, err := PushScope("plugin-generic-scope", ScopeTypeCustom) + if err != nil { + t.Fatalf("PushScope failed: %v", err) + } + if err := PopScope(handle); err != nil { + t.Fatalf("PopScope failed: %v", err) + } + if err := ClearPluginConfiguration(); err != nil { + t.Fatalf("ClearPluginConfiguration failed: %v", err) + } + if err := EmitEvent("plugin-mark-after-clear", WithEventData(json.RawMessage(`{"raw":true}`))); err != nil { + t.Fatalf("EmitEvent failed: %v", err) + } + if err := FlushSubscribers(); err != nil { + t.Fatalf("FlushSubscribers failed: %v", err) + } + mu.Lock() + defer mu.Unlock() + var phases []string + for _, event := range pluginEvents { + switch event.Name() { + case "plugin-mark": + if string(event.Data()) != `{"plugin":true}` { + t.Fatalf("plugin mark sanitizer did not run: %s", event.Data()) + } + case "plugin-mark-after-clear": + if string(event.Data()) != `{"raw":true}` { + t.Fatalf("plugin mark sanitizer was not rolled back: %s", event.Data()) + } + case "plugin-generic-scope": + phases = append(phases, string(event.Metadata())) + } + } + if len(phases) != 2 || phases[0] != `{"plugin_phase":"start"}` || phases[1] != `{"plugin_phase":"end"}` { + t.Fatalf("unexpected plugin scope sanitizer phases: %v", phases) + } } func TestPluginHelperConstructorsAndRegistryListing(t *testing.T) { @@ -424,6 +493,15 @@ func TestPluginFuncsAndClosedContextBranches(t *testing.T) { {"subscriber", func() error { return closed.RegisterSubscriber("subscriber", func(event Event) { _ = event }) }}, + {"mark sanitizer", func() error { + return closed.RegisterMarkSanitizeGuardrail("mark", 1, func(_ Event, fields EventSanitizeFields) EventSanitizeFields { return fields }) + }}, + {"scope start sanitizer", func() error { + return closed.RegisterScopeSanitizeStartGuardrail("scope_start", 1, func(_ Event, fields EventSanitizeFields) EventSanitizeFields { return fields }) + }}, + {"scope end sanitizer", func() error { + return closed.RegisterScopeSanitizeEndGuardrail("scope_end", 1, func(_ Event, fields EventSanitizeFields) EventSanitizeFields { return fields }) + }}, {"tool sanitize request", func() error { return closed.RegisterToolSanitizeRequestGuardrail("tool_sanitize_request", 1, func(name string, args json.RawMessage) json.RawMessage { return args }) }}, diff --git a/go/nemo_relay/callbacks.go b/go/nemo_relay/callbacks.go index 1ebab4914..f8069eda0 100644 --- a/go/nemo_relay/callbacks.go +++ b/go/nemo_relay/callbacks.go @@ -36,6 +36,7 @@ typedef char* (*NemoRelayLlmConditionalCb)(void* user_data, const FfiLLMRequest* typedef char* (*NemoRelayLlmExecFn)(void* user_data, const char* native_json); typedef char* (*NemoRelayLlmResponseFn)(void* user_data, const char* response_json); typedef void (*NemoRelayEventSubscriberFn)(void* user_data, const FfiEvent* event); +typedef char* (*NemoRelayEventSanitizeFn)(void* user_data, const FfiEvent* event, const char* fields_json); typedef struct FfiPluginContext FfiPluginContext; // Middleware chain next function types @@ -204,6 +205,22 @@ type FinalizerFunc func() string // the callback, so it is safe to retain the event after the callback returns. type EventSubscriberFunc func(event Event) +// EventSanitizeFields contains the observability fields an event sanitizer may replace. +// +// A sanitizer result replaces all three fields; it does not patch the input fields. +// Start with the provided fields and modify only the fields you intend to change. +// A zero-value json.RawMessage serializes as JSON null and clears its corresponding +// event field, so a newly constructed EventSanitizeFields clears every field left +// at its zero value. +type EventSanitizeFields struct { + Data json.RawMessage `json:"data"` + CategoryProfile json.RawMessage `json:"category_profile"` + Metadata json.RawMessage `json:"metadata"` +} + +// EventSanitizeFunc sanitizes mark or scope event observability fields. +type EventSanitizeFunc func(event Event, fields EventSanitizeFields) EventSanitizeFields + // CodecFunc is a bidirectional codec with decode and encode methods. // Decode receives the full LLM request (headers + content) as JSON and returns // the AnnotatedLLMRequest as JSON. Encode receives the annotated request JSON @@ -413,6 +430,22 @@ func goEventSubscriberTrampoline(userData unsafe.Pointer, event *C.FfiEvent) { fn(goEvent) } +//export goEventSanitizeTrampoline +func goEventSanitizeTrampoline(userData unsafe.Pointer, event *C.FfiEvent, fieldsJSON *C.char) *C.char { + fn := lookupClosure(userData).(EventSanitizeFunc) + var fields EventSanitizeFields + if err := json.Unmarshal([]byte(C.GoString(fieldsJSON)), &fields); err != nil { + setLastErrorMessage(err.Error()) + return nil + } + result, err := json.Marshal(fn(newEvent(event), fields)) + if err != nil { + setLastErrorMessage(err.Error()) + return nil + } + return C.CString(string(result)) +} + //export goFreeTrampoline func goFreeTrampoline(userData unsafe.Pointer) { unregisterClosure(userData) diff --git a/go/nemo_relay/event_sanitizers_test.go b/go/nemo_relay/event_sanitizers_test.go new file mode 100644 index 000000000..c0509536c --- /dev/null +++ b/go/nemo_relay/event_sanitizers_test.go @@ -0,0 +1,204 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package nemo_relay + +import ( + "encoding/json" + "sync" + "testing" +) + +func TestEventSanitizerRegistries(t *testing.T) { + var mu sync.Mutex + var events []Event + if err := RegisterSubscriber("go-event-sanitize-sub", func(event Event) { + mu.Lock() + events = append(events, event) + mu.Unlock() + }); err != nil { + t.Fatal(err) + } + defer DeregisterSubscriber("go-event-sanitize-sub") + + if err := RegisterMarkSanitizeGuardrail("go-mark-sanitize", 0, func(event Event, fields EventSanitizeFields) EventSanitizeFields { + if event.Name() != "checkpoint" { + t.Fatalf("unexpected event context: %s", event.Name()) + } + fields.Data = json.RawMessage(`{"safe":true}`) + fields.CategoryProfile = json.RawMessage(`{"subtype":"go.sanitized"}`) + fields.Metadata = json.RawMessage("null") + return fields + }); err != nil { + t.Fatal(err) + } + defer DeregisterMarkSanitizeGuardrail("go-mark-sanitize") + if err := RegisterScopeSanitizeStartGuardrail("go-scope-start", 0, func(_ Event, fields EventSanitizeFields) EventSanitizeFields { + fields.Metadata = json.RawMessage(`{"phase":"start"}`) + return fields + }); err != nil { + t.Fatal(err) + } + defer DeregisterScopeSanitizeStartGuardrail("go-scope-start") + if err := RegisterScopeSanitizeEndGuardrail("go-scope-end", 0, func(_ Event, fields EventSanitizeFields) EventSanitizeFields { + fields.Metadata = json.RawMessage(`{"phase":"end"}`) + return fields + }); err != nil { + t.Fatal(err) + } + defer DeregisterScopeSanitizeEndGuardrail("go-scope-end") + handle, err := PushScope("generic", ScopeTypeCustom) + if err != nil { + t.Fatal(err) + } + if err := PopScope(handle); err != nil { + t.Fatal(err) + } + + if err := EmitEvent("checkpoint", WithEventData(json.RawMessage(`{"secret":true}`)), WithEventMetadata(json.RawMessage(`{"secret":true}`))); err != nil { + t.Fatal(err) + } + if err := FlushSubscribers(); err != nil { + t.Fatal(err) + } + mu.Lock() + mark := events[len(events)-1] + mu.Unlock() + if string(mark.Data()) != `{"safe":true}` || string(mark.CategoryProfile()) != `{"subtype":"go.sanitized"}` || len(mark.Metadata()) != 0 { + t.Fatalf("unexpected sanitized fields: data=%s category_profile=%s metadata=%s", mark.Data(), mark.CategoryProfile(), mark.Metadata()) + } + var phases []string + for _, event := range events { + if event.Name() == "generic" { + phases = append(phases, string(event.Metadata())) + } + } + if len(phases) != 2 || phases[0] != `{"phase":"start"}` || phases[1] != `{"phase":"end"}` { + t.Fatalf("unexpected scope sanitizer phases: %v", phases) + } +} + +func TestScopeLocalEventSanitizerInheritanceAndCleanup(t *testing.T) { + var mu sync.Mutex + seen := map[string]json.RawMessage{} + if err := RegisterSubscriber("go-local-event-sub", func(event Event) { + mu.Lock() + seen[event.Name()+":"+event.ScopeCategory()] = append(json.RawMessage(nil), event.Data()...) + mu.Unlock() + }); err != nil { + t.Fatal(err) + } + defer DeregisterSubscriber("go-local-event-sub") + + owner, err := PushScope("owner", ScopeTypeAgent) + if err != nil { + t.Fatal(err) + } + if err := ScopeRegisterMarkSanitizeGuardrail(owner.UUID(), "go-local-mark", 0, func(_ Event, fields EventSanitizeFields) EventSanitizeFields { + fields.Data = json.RawMessage(`{"local":true}`) + return fields + }); err != nil { + t.Fatal(err) + } + if err := ScopeRegisterScopeSanitizeStartGuardrail(owner.UUID(), "go-local-start", 0, func(_ Event, fields EventSanitizeFields) EventSanitizeFields { + fields.Data = json.RawMessage(`{"local_phase":"start"}`) + return fields + }); err != nil { + t.Fatal(err) + } + if err := ScopeRegisterScopeSanitizeEndGuardrail(owner.UUID(), "go-local-end", 0, func(_ Event, fields EventSanitizeFields) EventSanitizeFields { + fields.Data = json.RawMessage(`{"local_phase":"end"}`) + return fields + }); err != nil { + t.Fatal(err) + } + if err := EmitEvent("inside", WithEventData(json.RawMessage(`{"raw":true}`))); err != nil { + t.Fatal(err) + } + child, err := PushScope("child", ScopeTypeFunction) + if err != nil { + t.Fatal(err) + } + if err := EmitEvent("inherited", WithEventData(json.RawMessage(`{"raw":true}`))); err != nil { + t.Fatal(err) + } + if err := PopScope(child); err != nil { + t.Fatal(err) + } + if err := PopScope(owner); err != nil { + t.Fatal(err) + } + if err := EmitEvent("outside", WithEventData(json.RawMessage(`{"raw":true}`))); err != nil { + t.Fatal(err) + } + if err := FlushSubscribers(); err != nil { + t.Fatal(err) + } + if string(seen["inside:"]) != `{"local":true}` || + string(seen["inherited:"]) != `{"local":true}` || + string(seen["outside:"]) != `{"raw":true}` || + string(seen["child:start"]) != `{"local_phase":"start"}` || + string(seen["child:end"]) != `{"local_phase":"end"}` { + t.Fatalf("unexpected scope-local results: %#v", seen) + } +} + +func TestEventSanitizerRegistrationErrorsReleaseCallbacks(t *testing.T) { + closureRegistryMu.Lock() + baseline := len(closureRegistry) + closureRegistryMu.Unlock() + passThrough := func(_ Event, fields EventSanitizeFields) EventSanitizeFields { return fields } + + if err := RegisterMarkSanitizeGuardrail("go-event-duplicate", 0, passThrough); err != nil { + t.Fatal(err) + } + if err := RegisterMarkSanitizeGuardrail("go-event-duplicate", 0, passThrough); err == nil { + t.Fatal("expected duplicate event sanitizer registration to fail") + } + closureRegistryMu.Lock() + afterDuplicate := len(closureRegistry) + closureRegistryMu.Unlock() + if afterDuplicate != baseline+1 { + t.Fatalf("duplicate registration leaked callback: baseline=%d current=%d", baseline, afterDuplicate) + } + if err := DeregisterMarkSanitizeGuardrail("go-event-duplicate"); err != nil { + t.Fatal(err) + } + if err := RegisterToolSanitizeRequestGuardrail("go-tool-duplicate", 0, func(_ string, args json.RawMessage) json.RawMessage { return args }); err != nil { + t.Fatal(err) + } + if err := RegisterToolSanitizeRequestGuardrail("go-tool-duplicate", 0, func(_ string, args json.RawMessage) json.RawMessage { return args }); err == nil { + t.Fatal("expected duplicate tool sanitizer registration to fail") + } + closureRegistryMu.Lock() + afterToolDuplicate := len(closureRegistry) + closureRegistryMu.Unlock() + if afterToolDuplicate != baseline+1 { + t.Fatalf("duplicate tool registration leaked callback: baseline=%d current=%d", baseline, afterToolDuplicate) + } + if err := DeregisterToolSanitizeRequestGuardrail("go-tool-duplicate"); err != nil { + t.Fatal(err) + } + + for name, register := range map[string]func() error{ + "mark": func() error { + return ScopeRegisterMarkSanitizeGuardrail("not-a-uuid", "go-invalid-mark", 0, passThrough) + }, + "scope start": func() error { + return ScopeRegisterScopeSanitizeStartGuardrail("not-a-uuid", "go-invalid-start", 0, passThrough) + }, + "scope end": func() error { + return ScopeRegisterScopeSanitizeEndGuardrail("not-a-uuid", "go-invalid-end", 0, passThrough) + }, + } { + if err := register(); err == nil { + t.Fatalf("expected invalid UUID for %s registration", name) + } + } + closureRegistryMu.Lock() + afterErrors := len(closureRegistry) + closureRegistryMu.Unlock() + if afterErrors != baseline { + t.Fatalf("failed registration leaked callbacks: baseline=%d current=%d", baseline, afterErrors) + } +} diff --git a/go/nemo_relay/guardrails/guardrails.go b/go/nemo_relay/guardrails/guardrails.go index 7fb212280..dc0388208 100644 --- a/go/nemo_relay/guardrails/guardrails.go +++ b/go/nemo_relay/guardrails/guardrails.go @@ -34,6 +34,66 @@ import ( "github.com/NVIDIA/NeMo-Relay/go/nemo_relay" ) +// RegisterMarkSanitize registers a global mark event sanitizer. +func RegisterMarkSanitize(name string, priority int32, fn nemo_relay.EventSanitizeFunc) error { + return nemo_relay.RegisterMarkSanitizeGuardrail(name, priority, fn) +} + +// DeregisterMarkSanitize removes a global mark event sanitizer. +func DeregisterMarkSanitize(name string) error { + return nemo_relay.DeregisterMarkSanitizeGuardrail(name) +} + +// RegisterScopeSanitizeStart registers a global scope-start event sanitizer. +func RegisterScopeSanitizeStart(name string, priority int32, fn nemo_relay.EventSanitizeFunc) error { + return nemo_relay.RegisterScopeSanitizeStartGuardrail(name, priority, fn) +} + +// DeregisterScopeSanitizeStart removes a global scope-start event sanitizer. +func DeregisterScopeSanitizeStart(name string) error { + return nemo_relay.DeregisterScopeSanitizeStartGuardrail(name) +} + +// RegisterScopeSanitizeEnd registers a global scope-end event sanitizer. +func RegisterScopeSanitizeEnd(name string, priority int32, fn nemo_relay.EventSanitizeFunc) error { + return nemo_relay.RegisterScopeSanitizeEndGuardrail(name, priority, fn) +} + +// DeregisterScopeSanitizeEnd removes a global scope-end event sanitizer. +func DeregisterScopeSanitizeEnd(name string) error { + return nemo_relay.DeregisterScopeSanitizeEndGuardrail(name) +} + +// ScopeRegisterMarkSanitize registers a scope-local mark event sanitizer. +func ScopeRegisterMarkSanitize(scopeUUID, name string, priority int32, fn nemo_relay.EventSanitizeFunc) error { + return nemo_relay.ScopeRegisterMarkSanitizeGuardrail(scopeUUID, name, priority, fn) +} + +// ScopeDeregisterMarkSanitize removes a scope-local mark event sanitizer. +func ScopeDeregisterMarkSanitize(scopeUUID, name string) error { + return nemo_relay.ScopeDeregisterMarkSanitizeGuardrail(scopeUUID, name) +} + +// ScopeRegisterScopeSanitizeStart registers a scope-local scope-start sanitizer. +func ScopeRegisterScopeSanitizeStart(scopeUUID, name string, priority int32, fn nemo_relay.EventSanitizeFunc) error { + return nemo_relay.ScopeRegisterScopeSanitizeStartGuardrail(scopeUUID, name, priority, fn) +} + +// ScopeDeregisterScopeSanitizeStart removes a scope-local scope-start sanitizer. +func ScopeDeregisterScopeSanitizeStart(scopeUUID, name string) error { + return nemo_relay.ScopeDeregisterScopeSanitizeStartGuardrail(scopeUUID, name) +} + +// ScopeRegisterScopeSanitizeEnd registers a scope-local scope-end sanitizer. +func ScopeRegisterScopeSanitizeEnd(scopeUUID, name string, priority int32, fn nemo_relay.EventSanitizeFunc) error { + return nemo_relay.ScopeRegisterScopeSanitizeEndGuardrail(scopeUUID, name, priority, fn) +} + +// ScopeDeregisterScopeSanitizeEnd removes a scope-local scope-end sanitizer. +func ScopeDeregisterScopeSanitizeEnd(scopeUUID, name string) error { + return nemo_relay.ScopeDeregisterScopeSanitizeEndGuardrail(scopeUUID, name) +} + // --- Tool Sanitize Request --- // RegisterToolSanitizeRequest registers a guardrail that sanitizes tool request diff --git a/go/nemo_relay/guardrails/guardrails_test.go b/go/nemo_relay/guardrails/guardrails_test.go index 1e8ca7839..002906950 100644 --- a/go/nemo_relay/guardrails/guardrails_test.go +++ b/go/nemo_relay/guardrails/guardrails_test.go @@ -249,6 +249,22 @@ func runScopeLocalLLMGuardrailShorthandChecks(t *testing.T, scopeUUID string) { } func TestGuardrailShorthandsGlobal(t *testing.T) { + eventFn := func(_ nemo_relay.Event, fields nemo_relay.EventSanitizeFields) nemo_relay.EventSanitizeFields { + return fields + } + if err := guardrails.RegisterMarkSanitize("guardrails_mark", 0, eventFn); err != nil { + t.Fatal(err) + } + if err := guardrails.RegisterScopeSanitizeStart("guardrails_start", 0, eventFn); err != nil { + t.Fatal(err) + } + if err := guardrails.RegisterScopeSanitizeEnd("guardrails_end", 0, eventFn); err != nil { + t.Fatal(err) + } + defer guardrails.DeregisterMarkSanitize("guardrails_mark") + defer guardrails.DeregisterScopeSanitizeStart("guardrails_start") + defer guardrails.DeregisterScopeSanitizeEnd("guardrails_end") + _ = nemo_relay.EmitEvent("guardrails-mark") toolEventOutput, cleanupTool := captureEndEventOutput(t, "guardrails_tool_events", "guardrails_tool") defer cleanupTool() runGlobalToolGuardrailShorthandChecks(t, toolEventOutput) @@ -273,6 +289,22 @@ func TestGuardrailShorthandsScopeLocal(t *testing.T) { defer nemo_relay.PopScope(handle) scopeUUID := handle.UUID() + eventFn := func(_ nemo_relay.Event, fields nemo_relay.EventSanitizeFields) nemo_relay.EventSanitizeFields { + return fields + } + if err := guardrails.ScopeRegisterMarkSanitize(scopeUUID, "guardrails_scope_mark", 0, eventFn); err != nil { + t.Fatal(err) + } + if err := guardrails.ScopeRegisterScopeSanitizeStart(scopeUUID, "guardrails_scope_start", 0, eventFn); err != nil { + t.Fatal(err) + } + if err := guardrails.ScopeRegisterScopeSanitizeEnd(scopeUUID, "guardrails_scope_end", 0, eventFn); err != nil { + t.Fatal(err) + } + defer guardrails.ScopeDeregisterMarkSanitize(scopeUUID, "guardrails_scope_mark") + defer guardrails.ScopeDeregisterScopeSanitizeStart(scopeUUID, "guardrails_scope_start") + defer guardrails.ScopeDeregisterScopeSanitizeEnd(scopeUUID, "guardrails_scope_end") + _ = nemo_relay.EmitEvent("guardrails-scope-mark") runScopeLocalToolGuardrailShorthandChecks(t, scopeUUID) runScopeLocalLLMGuardrailShorthandChecks(t, scopeUUID) }) diff --git a/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index 87ab55291..b09718ff5 100644 --- a/go/nemo_relay/nemo_relay.go +++ b/go/nemo_relay/nemo_relay.go @@ -161,11 +161,24 @@ extern int32_t nemo_relay_deregister_llm_stream_execution_intercept(const char* // Subscribers typedef void (*NemoRelayEventSubscriberFn)(void* user_data, const FfiEvent* event); +typedef char* (*NemoRelayEventSanitizeFn)(void* user_data, const FfiEvent* event, const char* fields_json); extern int32_t nemo_relay_register_subscriber(const char* name, NemoRelayEventSubscriberFn cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_deregister_subscriber(const char* name); extern int32_t nemo_relay_flush_subscribers(void); +extern int32_t nemo_relay_register_mark_sanitize_guardrail(const char* name, int32_t priority, NemoRelayEventSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_deregister_mark_sanitize_guardrail(const char* name); +extern int32_t nemo_relay_register_scope_sanitize_start_guardrail(const char* name, int32_t priority, NemoRelayEventSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_deregister_scope_sanitize_start_guardrail(const char* name); +extern int32_t nemo_relay_register_scope_sanitize_end_guardrail(const char* name, int32_t priority, NemoRelayEventSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_deregister_scope_sanitize_end_guardrail(const char* name); // Scope-local tool guardrails +extern int32_t nemo_relay_scope_register_mark_sanitize_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayEventSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_scope_deregister_mark_sanitize_guardrail(const char* scope_uuid, const char* name); +extern int32_t nemo_relay_scope_register_scope_sanitize_start_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayEventSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_scope_deregister_scope_sanitize_start_guardrail(const char* scope_uuid, const char* name); +extern int32_t nemo_relay_scope_register_scope_sanitize_end_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayEventSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_scope_deregister_scope_sanitize_end_guardrail(const char* scope_uuid, const char* name); extern int32_t nemo_relay_scope_register_tool_sanitize_request_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayToolSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_scope_deregister_tool_sanitize_request_guardrail(const char* scope_uuid, const char* name); extern int32_t nemo_relay_scope_register_tool_sanitize_response_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayToolSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); @@ -254,6 +267,7 @@ extern void nemo_relay_openinference_subscriber_free(void*); // Go trampoline forward declarations (defined via //export in callbacks.go) extern char* goToolSanitizeTrampoline(void*, const char*, const char*); +extern char* goEventSanitizeTrampoline(void*, const FfiEvent*, const char*); extern char* goToolConditionalTrampoline(void*, const char*, const char*); extern char* goToolExecTrampoline(void*, const char*); extern void goEventSubscriberTrampoline(void*, const FfiEvent*); @@ -1190,6 +1204,58 @@ func LlmStreamCallExecute(name string, request interface{}, fn LLMExecutionFunc, // Guardrail/Intercept registration (Tool) // --------------------------------------------------------------------------- +func registerEventSanitizer(name string, priority int32, fn EventSanitizeFunc, kind int) error { + id := registerClosure(fn) + cName := C.CString(name) + defer C.free(unsafe.Pointer(cName)) + var status C.int32_t + switch kind { + case 0: + status = C.nemo_relay_register_mark_sanitize_guardrail(cName, C.int32_t(priority), C.NemoRelayEventSanitizeFn(C.goEventSanitizeTrampoline), id, C.NemoRelayFreeFn(C.goFreeTrampoline)) + case 1: + status = C.nemo_relay_register_scope_sanitize_start_guardrail(cName, C.int32_t(priority), C.NemoRelayEventSanitizeFn(C.goEventSanitizeTrampoline), id, C.NemoRelayFreeFn(C.goFreeTrampoline)) + default: + status = C.nemo_relay_register_scope_sanitize_end_guardrail(cName, C.int32_t(priority), C.NemoRelayEventSanitizeFn(C.goEventSanitizeTrampoline), id, C.NemoRelayFreeFn(C.goFreeTrampoline)) + } + return checkStatus(status) +} + +// RegisterMarkSanitizeGuardrail registers a global mark event sanitizer. +func RegisterMarkSanitizeGuardrail(name string, priority int32, fn EventSanitizeFunc) error { + return registerEventSanitizer(name, priority, fn, 0) +} + +// DeregisterMarkSanitizeGuardrail removes a global mark event sanitizer. +func DeregisterMarkSanitizeGuardrail(name string) error { + cName := C.CString(name) + defer C.free(unsafe.Pointer(cName)) + return checkStatus(C.nemo_relay_deregister_mark_sanitize_guardrail(cName)) +} + +// RegisterScopeSanitizeStartGuardrail registers a global scope-start event sanitizer. +func RegisterScopeSanitizeStartGuardrail(name string, priority int32, fn EventSanitizeFunc) error { + return registerEventSanitizer(name, priority, fn, 1) +} + +// DeregisterScopeSanitizeStartGuardrail removes a global scope-start event sanitizer. +func DeregisterScopeSanitizeStartGuardrail(name string) error { + cName := C.CString(name) + defer C.free(unsafe.Pointer(cName)) + return checkStatus(C.nemo_relay_deregister_scope_sanitize_start_guardrail(cName)) +} + +// RegisterScopeSanitizeEndGuardrail registers a global scope-end event sanitizer. +func RegisterScopeSanitizeEndGuardrail(name string, priority int32, fn EventSanitizeFunc) error { + return registerEventSanitizer(name, priority, fn, 2) +} + +// DeregisterScopeSanitizeEndGuardrail removes a global scope-end event sanitizer. +func DeregisterScopeSanitizeEndGuardrail(name string) error { + cName := C.CString(name) + defer C.free(unsafe.Pointer(cName)) + return checkStatus(C.nemo_relay_deregister_scope_sanitize_end_guardrail(cName)) +} + // RegisterToolSanitizeRequestGuardrail registers a guardrail that sanitizes // tool request arguments before they are passed to the tool. The callback // receives the tool name and arguments JSON and must return the (possibly @@ -2073,6 +2139,68 @@ func (s *OpenInferenceSubscriber) Close() { // Scope-local guardrail/intercept registration (Tool) // --------------------------------------------------------------------------- +func registerScopeEventSanitizer(scopeUUID, name string, priority int32, fn EventSanitizeFunc, kind int) error { + id := registerClosure(fn) + cScopeUUID := C.CString(scopeUUID) + defer C.free(unsafe.Pointer(cScopeUUID)) + cName := C.CString(name) + defer C.free(unsafe.Pointer(cName)) + callback := C.NemoRelayEventSanitizeFn(C.goEventSanitizeTrampoline) + free := C.NemoRelayFreeFn(C.goFreeTrampoline) + var status C.int32_t + switch kind { + case 0: + status = C.nemo_relay_scope_register_mark_sanitize_guardrail(cScopeUUID, cName, C.int32_t(priority), callback, id, free) + case 1: + status = C.nemo_relay_scope_register_scope_sanitize_start_guardrail(cScopeUUID, cName, C.int32_t(priority), callback, id, free) + default: + status = C.nemo_relay_scope_register_scope_sanitize_end_guardrail(cScopeUUID, cName, C.int32_t(priority), callback, id, free) + } + return checkStatus(status) +} + +// ScopeRegisterMarkSanitizeGuardrail registers a scope-local mark sanitizer. +func ScopeRegisterMarkSanitizeGuardrail(scopeUUID, name string, priority int32, fn EventSanitizeFunc) error { + return registerScopeEventSanitizer(scopeUUID, name, priority, fn, 0) +} + +// ScopeDeregisterMarkSanitizeGuardrail removes a scope-local mark sanitizer. +func ScopeDeregisterMarkSanitizeGuardrail(scopeUUID, name string) error { + cScopeUUID := C.CString(scopeUUID) + defer C.free(unsafe.Pointer(cScopeUUID)) + cName := C.CString(name) + defer C.free(unsafe.Pointer(cName)) + return checkStatus(C.nemo_relay_scope_deregister_mark_sanitize_guardrail(cScopeUUID, cName)) +} + +// ScopeRegisterScopeSanitizeStartGuardrail registers a scope-local scope-start sanitizer. +func ScopeRegisterScopeSanitizeStartGuardrail(scopeUUID, name string, priority int32, fn EventSanitizeFunc) error { + return registerScopeEventSanitizer(scopeUUID, name, priority, fn, 1) +} + +// ScopeDeregisterScopeSanitizeStartGuardrail removes a scope-local scope-start sanitizer. +func ScopeDeregisterScopeSanitizeStartGuardrail(scopeUUID, name string) error { + cScopeUUID := C.CString(scopeUUID) + defer C.free(unsafe.Pointer(cScopeUUID)) + cName := C.CString(name) + defer C.free(unsafe.Pointer(cName)) + return checkStatus(C.nemo_relay_scope_deregister_scope_sanitize_start_guardrail(cScopeUUID, cName)) +} + +// ScopeRegisterScopeSanitizeEndGuardrail registers a scope-local scope-end sanitizer. +func ScopeRegisterScopeSanitizeEndGuardrail(scopeUUID, name string, priority int32, fn EventSanitizeFunc) error { + return registerScopeEventSanitizer(scopeUUID, name, priority, fn, 2) +} + +// ScopeDeregisterScopeSanitizeEndGuardrail removes a scope-local scope-end sanitizer. +func ScopeDeregisterScopeSanitizeEndGuardrail(scopeUUID, name string) error { + cScopeUUID := C.CString(scopeUUID) + defer C.free(unsafe.Pointer(cScopeUUID)) + cName := C.CString(name) + defer C.free(unsafe.Pointer(cName)) + return checkStatus(C.nemo_relay_scope_deregister_scope_sanitize_end_guardrail(cScopeUUID, cName)) +} + // ScopeRegisterToolSanitizeRequestGuardrail registers a scope-local guardrail // that sanitizes tool request arguments. The guardrail is scoped to the given // scope UUID and does not affect other scopes. diff --git a/go/nemo_relay/pii_redaction.go b/go/nemo_relay/pii_redaction.go index 691e6a7aa..871a83b60 100644 --- a/go/nemo_relay/pii_redaction.go +++ b/go/nemo_relay/pii_redaction.go @@ -33,6 +33,7 @@ type PiiRedactionConfig struct { Mode string `json:"mode,omitempty"` Input bool `json:"input"` Output bool `json:"output"` + Mark bool `json:"mark"` ToolInput bool `json:"tool_input"` ToolOutput bool `json:"tool_output"` Priority int32 `json:"priority,omitempty"` @@ -56,6 +57,7 @@ func NewPiiRedactionConfig() PiiRedactionConfig { Mode: "builtin", Input: true, Output: true, + Mark: true, ToolInput: true, ToolOutput: true, Priority: 100, diff --git a/go/nemo_relay/pii_redaction_test.go b/go/nemo_relay/pii_redaction_test.go index c0fd4dba9..5b5728adc 100644 --- a/go/nemo_relay/pii_redaction_test.go +++ b/go/nemo_relay/pii_redaction_test.go @@ -7,7 +7,7 @@ import "testing" func TestPiiRedactionConfigHelpers(t *testing.T) { config := NewPiiRedactionConfig() - if config.Version != 1 || config.Mode != "builtin" || !config.Input || !config.Output || !config.ToolInput || !config.ToolOutput || config.Priority != 100 { + if config.Version != 1 || config.Mode != "builtin" || !config.Input || !config.Output || !config.Mark || !config.ToolInput || !config.ToolOutput || config.Priority != 100 { t.Fatalf("unexpected PII redaction defaults: %#v", config) } if config.Builtin == nil || config.Builtin.Action != "remove" || len(config.Builtin.TargetPaths) != 0 { @@ -30,6 +30,9 @@ func TestPiiRedactionConfigHelpers(t *testing.T) { if component.Config["mode"] != "builtin" || component.Config["priority"] != float64(100) { t.Fatalf("unexpected serialized config: %#v", component.Config) } + if component.Config["mark"] != true { + t.Fatalf("expected mark redaction to be enabled: %#v", component.Config) + } serializedBuiltin, ok := component.Config["builtin"].(map[string]any) if !ok { t.Fatalf("expected serialized builtin object, got %#v", component.Config["builtin"]) diff --git a/go/nemo_relay/plugin.go b/go/nemo_relay/plugin.go index af39171b6..1ec43b80e 100644 --- a/go/nemo_relay/plugin.go +++ b/go/nemo_relay/plugin.go @@ -13,6 +13,7 @@ typedef void (*NemoRelayFreeFn)(void* user_data); typedef char* (*NemoRelayPluginValidateCb)(void* user_data, const char* plugin_config_json); typedef int32_t (*NemoRelayPluginRegisterCb)(void* user_data, const char* plugin_config_json, FfiPluginContext* ctx); typedef void (*NemoRelayEventSubscriberFn)(void* user_data, const void* event); +typedef char* (*NemoRelayEventSanitizeFn)(void* user_data, const void* event, const char* fields_json); typedef char* (*NemoRelayToolSanitizeFn)(void* user_data, const char* name, const char* args_json); typedef char* (*NemoRelayToolConditionalFn)(void* user_data, const char* name, const char* args_json); typedef void* (*NemoRelayLlmRequestCb)(void* user_data, const void* request); @@ -34,6 +35,9 @@ extern int32_t nemo_relay_deregister_plugin(const char* plugin_kind); extern void nemo_relay_string_free(char* ptr); extern int32_t nemo_relay_plugin_context_register_subscriber(FfiPluginContext* ctx, const char* name, NemoRelayEventSubscriberFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_plugin_context_register_mark_sanitize_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayEventSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_plugin_context_register_scope_sanitize_start_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayEventSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_plugin_context_register_scope_sanitize_end_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayEventSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_tool_sanitize_request_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayToolSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_tool_sanitize_response_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayToolSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_tool_conditional_execution_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayToolConditionalFn cb, void* user_data, NemoRelayFreeFn free_fn); @@ -49,6 +53,7 @@ extern int32_t nemo_relay_plugin_context_register_tool_execution_intercept(FfiPl extern char* goPluginValidateTrampoline(void*, const char*); extern int32_t goPluginRegisterTrampoline(void*, const char*, FfiPluginContext*); extern void goEventSubscriberTrampoline(void*, const void*); +extern char* goEventSanitizeTrampoline(void*, const void*, const char*); extern void goFreeTrampoline(void*); extern char* goToolSanitizeTrampoline(void*, const char*, const char*); extern char* goToolConditionalTrampoline(void*, const char*, const char*); @@ -342,6 +347,42 @@ func (ctx *PluginContext) RegisterSubscriber(name string, fn EventSubscriberFunc )) } +func (ctx *PluginContext) registerEventSanitizer(name string, priority int32, fn EventSanitizeFunc, surface int) error { + if ctx == nil || ctx.ptr == nil { + return errors.New(errPluginContextClosed) + } + cName := C.CString(name) + defer C.free(unsafe.Pointer(cName)) + userData := registerClosure(fn) + callback := C.NemoRelayEventSanitizeFn(C.goEventSanitizeTrampoline) + free := C.NemoRelayFreeFn(C.goFreeTrampoline) + var status C.int32_t + switch surface { + case 0: + status = C.nemo_relay_plugin_context_register_mark_sanitize_guardrail(ctx.ptr, cName, C.int32_t(priority), callback, userData, free) + case 1: + status = C.nemo_relay_plugin_context_register_scope_sanitize_start_guardrail(ctx.ptr, cName, C.int32_t(priority), callback, userData, free) + default: + status = C.nemo_relay_plugin_context_register_scope_sanitize_end_guardrail(ctx.ptr, cName, C.int32_t(priority), callback, userData, free) + } + return checkStatus(status) +} + +// RegisterMarkSanitizeGuardrail registers a mark event sanitizer for this component. +func (ctx *PluginContext) RegisterMarkSanitizeGuardrail(name string, priority int32, fn EventSanitizeFunc) error { + return ctx.registerEventSanitizer(name, priority, fn, 0) +} + +// RegisterScopeSanitizeStartGuardrail registers a scope-start sanitizer for this component. +func (ctx *PluginContext) RegisterScopeSanitizeStartGuardrail(name string, priority int32, fn EventSanitizeFunc) error { + return ctx.registerEventSanitizer(name, priority, fn, 1) +} + +// RegisterScopeSanitizeEndGuardrail registers a scope-end sanitizer for this component. +func (ctx *PluginContext) RegisterScopeSanitizeEndGuardrail(name string, priority int32, fn EventSanitizeFunc) error { + return ctx.registerEventSanitizer(name, priority, fn, 2) +} + // RegisterToolSanitizeRequestGuardrail registers a tool sanitize-request guardrail for this component. func (ctx *PluginContext) RegisterToolSanitizeRequestGuardrail(name string, priority int32, fn ToolSanitizeFunc) error { if ctx == nil || ctx.ptr == nil {