diff --git a/crates/cli/src/server/mod.rs b/crates/cli/src/server/mod.rs index a263f0a54..2bfc568f9 100644 --- a/crates/cli/src/server/mod.rs +++ b/crates/cli/src/server/mod.rs @@ -25,6 +25,7 @@ use nemo_relay::plugin::dynamic::{ }; use nemo_relay::plugin::{ PluginComponentSpec, PluginConfig, clear_plugin_configuration, initialize_plugins_exact, + initialize_plugins_exact_with_inference_providers, }; use nemo_relay_adaptive::plugin_component::register_adaptive_component; use nemo_relay_pii_redaction::component::register_pii_redaction_component; @@ -1116,7 +1117,11 @@ impl PluginActivation { CliError::Config(format!("worker plugin load failed: {error}")) })?) }; - initialize_plugins_exact(plugin_config) + let inference_providers = worker + .as_ref() + .map(WorkerPluginActivation::inference_providers) + .unwrap_or_default(); + initialize_plugins_exact_with_inference_providers(plugin_config, inference_providers) .await .map_err(|error| CliError::Config(format!("plugin activation failed: {error}")))?; Ok(Self { diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 18c9f5df1..6505e1e5f 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -69,6 +69,9 @@ mod registry; pub mod shared_runtime; pub mod stream; +#[cfg(test)] +#[path = "../tests/unit/inference_provider_tests.rs"] +mod inference_provider_tests; #[cfg(test)] #[path = "../tests/unit/types_tests.rs"] mod types_tests; diff --git a/crates/core/src/plugin.rs b/crates/core/src/plugin.rs index ff7bc6c79..0a8305d8a 100644 --- a/crates/core/src/plugin.rs +++ b/crates/core/src/plugin.rs @@ -48,6 +48,12 @@ pub use nemo_relay_types::plugin::{ConfigDiagnostic, DiagnosticLevel}; pub mod dynamic; pub use dynamic::*; +mod inference; +#[doc(hidden)] +pub use inference::{ + InferenceProvider, InferenceProviderDescriptor, InferenceProviderFn, + InferenceProviderRegistration, InferenceProviderRegistry, +}; type PluginMap = HashMap; @@ -342,6 +348,7 @@ impl PluginRegistration { pub struct PluginRegistrationContext { registrations: Vec, namespace: Option, + inference_providers: InferenceProviderRegistry, } impl PluginRegistrationContext { @@ -355,9 +362,33 @@ impl PluginRegistrationContext { Self { registrations: vec![], namespace: Some(namespace.into()), + inference_providers: InferenceProviderRegistry::default(), } } + /// Creates a registration context backed by host-owned inference providers. + #[doc(hidden)] + pub fn with_inference_providers( + namespace: Option, + inference_providers: InferenceProviderRegistry, + ) -> Self { + Self { + registrations: Vec::new(), + namespace, + inference_providers, + } + } + + /// Resolves an inference provider implementing the required contract. + #[doc(hidden)] + pub fn inference_provider( + &self, + name: &str, + expected_contract: &str, + ) -> Result { + self.inference_providers.resolve(name, expected_contract) + } + /// Returns the runtime-qualified name for a plugin-local registration. /// /// Plugin handlers should pass stable component-local names such as @@ -1309,12 +1340,23 @@ pub fn plugin_config_schema() -> Json { /// is removed before the new configuration is activated. #[doc(hidden)] pub async fn initialize_plugins_exact(config: PluginConfig) -> Result { + initialize_plugins_exact_with_inference_providers(config, InferenceProviderRegistry::default()) + .await +} + +/// Configures plugin components with host-owned inference providers. +#[doc(hidden)] +pub async fn initialize_plugins_exact_with_inference_providers( + config: PluginConfig, + inference_providers: InferenceProviderRegistry, +) -> Result { run_owned_plugin_mutation("plugin initialization", move || async move { let lease = LegacyPluginMutationLease::acquire()?; let rollback_failures = Arc::new(Mutex::new(Vec::new())); let initialization = tokio::spawn(initialize_plugins_exact_inner( config, Some(Arc::clone(&rollback_failures)), + inference_providers, )) .await .map_err(|error| { @@ -1431,14 +1473,16 @@ pub(crate) async fn initialize_plugins_exact_for_host( config: PluginConfig, owner_id: u64, rollback_failures: Arc>>, + inference_providers: InferenceProviderRegistry, ) -> Result { verify_plugin_host_owner(owner_id)?; - initialize_plugins_exact_inner(config, Some(rollback_failures)).await + initialize_plugins_exact_inner(config, Some(rollback_failures), inference_providers).await } async fn initialize_plugins_exact_inner( config: PluginConfig, rollback_failures: Option>>>, + inference_providers: InferenceProviderRegistry, ) -> Result { let enabled_component_count = config .components @@ -1475,11 +1519,17 @@ async fn initialize_plugins_exact_inner( match initialize_plugin_components_catching_panics( config.clone(), rollback_failures.clone(), + inference_providers.clone(), ) .await { Ok(registrations) => { - store_active_plugin_configuration(config, report.clone(), registrations)?; + store_active_plugin_configuration( + config, + report.clone(), + registrations, + inference_providers, + )?; log::info!( target: "nemo_relay.plugin", event = "plugin_configuration_replaced", @@ -1491,6 +1541,7 @@ async fn initialize_plugins_exact_inner( Err(err) => match initialize_plugin_components_catching_panics( previous_state.config.clone(), rollback_failures.clone(), + previous_state.inference_providers.clone(), ) .await { @@ -1500,6 +1551,7 @@ async fn initialize_plugins_exact_inner( previous_state.config, previous_report, registrations, + previous_state.inference_providers, )?; log::warn!( target: "nemo_relay.plugin", @@ -1523,9 +1575,18 @@ async fn initialize_plugins_exact_inner( }, } } else { - let registrations = - initialize_plugin_components_catching_panics(config.clone(), rollback_failures).await?; - store_active_plugin_configuration(config, report.clone(), registrations)?; + let registrations = initialize_plugin_components_catching_panics( + config.clone(), + rollback_failures, + inference_providers.clone(), + ) + .await?; + store_active_plugin_configuration( + config, + report.clone(), + registrations, + inference_providers, + )?; log::info!( target: "nemo_relay.plugin", event = "plugin_configuration_activated", @@ -1539,14 +1600,17 @@ async fn initialize_plugins_exact_inner( async fn initialize_plugin_components_catching_panics( config: PluginConfig, rollback_failures: Option>>>, + inference_providers: InferenceProviderRegistry, ) -> Result> { - tokio::spawn(async move { initialize_plugin_components(&config, rollback_failures).await }) - .await - .map_err(|error| { - PluginError::Internal(format!( - "plugin component initialization task failed: {error}" - )) - })? + tokio::spawn(async move { + initialize_plugin_components(&config, rollback_failures, inference_providers).await + }) + .await + .map_err(|error| { + PluginError::Internal(format!( + "plugin component initialization task failed: {error}" + )) + })? } /// Validates and activates `config` layered on top of the discovered @@ -1559,6 +1623,16 @@ pub async fn initialize_plugins(config: PluginConfig) -> Result { initialize_plugins_exact(config).await } +/// Resolves discovered configuration and activates it with host inference providers. +#[doc(hidden)] +pub async fn initialize_plugins_with_inference_providers( + config: PluginConfig, + inference_providers: InferenceProviderRegistry, +) -> Result { + let config = resolve_plugin_config(config)?; + initialize_plugins_exact_with_inference_providers(config, inference_providers).await +} + /// Layers `config` over the default discovered `plugins.toml` files. /// /// This is crate-visible so owned dynamic-plugin activation can use the same @@ -1955,11 +2029,13 @@ struct ActivePluginConfiguration { config: PluginConfig, report: ConfigReport, registrations: Vec, + inference_providers: InferenceProviderRegistry, } async fn initialize_plugin_components( config: &PluginConfig, rollback_failures: Option>>>, + inference_providers: InferenceProviderRegistry, ) -> Result> { ensure_builtin_plugins_registered()?; let totals = plugin_component_totals(config); @@ -1988,8 +2064,11 @@ async fn initialize_plugin_components( totals.get(component.kind.as_str()).copied().unwrap_or(1), ); - let mut pending = - PendingPluginRegistrationContext::new(namespace, rollback_failures.clone()); + let mut pending = PendingPluginRegistrationContext::new( + namespace, + rollback_failures.clone(), + inference_providers.clone(), + ); plugin .register(&component.config, &mut pending.context) .await?; @@ -2034,9 +2113,16 @@ struct PendingPluginRegistrationContext { } impl PendingPluginRegistrationContext { - fn new(namespace: String, rollback_failures: Option>>>) -> Self { + fn new( + namespace: String, + rollback_failures: Option>>>, + inference_providers: InferenceProviderRegistry, + ) -> Self { Self { - context: PluginRegistrationContext::with_namespace(namespace), + context: PluginRegistrationContext::with_inference_providers( + Some(namespace), + inference_providers, + ), rollback_failures, } } @@ -2071,6 +2157,7 @@ fn store_active_plugin_configuration( config: PluginConfig, report: ConfigReport, registrations: Vec, + inference_providers: InferenceProviderRegistry, ) -> Result<()> { let mut guard = ACTIVE_PLUGIN_CONFIGURATION.lock().map_err(|err| { PluginError::Internal(format!("active plugin configuration lock poisoned: {err}")) @@ -2079,6 +2166,7 @@ fn store_active_plugin_configuration( config, report, registrations, + inference_providers, }); Ok(()) } diff --git a/crates/core/src/plugin/dynamic/host.rs b/crates/core/src/plugin/dynamic/host.rs index 8b3f3e2df..fe55a803d 100644 --- a/crates/core/src/plugin/dynamic/host.rs +++ b/crates/core/src/plugin/dynamic/host.rs @@ -16,8 +16,8 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value as Json}; use crate::plugin::{ - ConfigReport, PluginComponentSpec, PluginConfig, PluginHostLease, Result, - acquire_plugin_host_lease, clear_plugin_configuration_for_host, + ConfigReport, InferenceProviderRegistry, PluginComponentSpec, PluginConfig, PluginHostLease, + Result, acquire_plugin_host_lease, clear_plugin_configuration_for_host, ensure_builtin_plugins_registered, initialize_plugins_exact_for_host, resolve_plugin_config, run_owned_plugin_mutation, }; @@ -186,10 +186,18 @@ impl PluginHostActivation { ); let rollback_failures = Arc::new(Mutex::new(Vec::new())); let owner_id = claim.owner_id(); + #[cfg(feature = "worker-grpc")] + let inference_providers = worker + .as_ref() + .map(WorkerPluginActivation::inference_providers) + .unwrap_or_else(InferenceProviderRegistry::default); + #[cfg(not(feature = "worker-grpc"))] + let inference_providers = InferenceProviderRegistry::default(); let initialization = tokio::spawn(initialize_plugins_exact_for_host( config, owner_id, Arc::clone(&rollback_failures), + inference_providers, )) .await .map_err(|error| { @@ -262,6 +270,16 @@ impl PluginHostActivation { self.active } + /// Returns the inference providers owned by this activation. + #[doc(hidden)] + pub fn inference_providers(&self) -> InferenceProviderRegistry { + #[cfg(feature = "worker-grpc")] + if let Some(worker) = &self.worker { + return worker.inference_providers(); + } + InferenceProviderRegistry::default() + } + /// Clear registered callbacks before unloading libraries and workers. pub fn clear(mut self) -> Result<()> { self.clear_inner() @@ -300,6 +318,7 @@ impl PluginHostActivation { #[cfg(feature = "worker-grpc")] if let Some(worker) = &mut self.worker { runtime_outcome.merge(worker.deregister_plugin_kinds_checked()); + runtime_outcome.merge(worker.deregister_inference_providers_checked()); } // A worker cannot be stopped while its registry adapter might still be diff --git a/crates/core/src/plugin/dynamic/worker.rs b/crates/core/src/plugin/dynamic/worker.rs index 79e2e70ce..5bb5e2cbc 100644 --- a/crates/core/src/plugin/dynamic/worker.rs +++ b/crates/core/src/plugin/dynamic/worker.rs @@ -72,8 +72,9 @@ use crate::codec::request::{ANNOTATED_LLM_REQUEST_SCHEMA, AnnotatedLlmRequest}; use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::{FlowError, Result as FlowResult}; use crate::plugin::{ - ConfigDiagnostic, DiagnosticLevel, Plugin, PluginError, PluginRegistrationContext, - deregister_plugin_registration_checked, register_plugin_tracked, + ConfigDiagnostic, DiagnosticLevel, InferenceProviderDescriptor, InferenceProviderRegistration, + InferenceProviderRegistry, Plugin, PluginDeregistrationOutcome, PluginError, + PluginRegistrationContext, deregister_plugin_registration_checked, register_plugin_tracked, }; use super::{ @@ -129,6 +130,8 @@ pub struct WorkerPluginLoadSpec { pub struct WorkerPluginActivation { plugins: Vec>, plugin_registrations: Vec<(String, u64)>, + inference_providers: InferenceProviderRegistry, + inference_provider_registrations: Vec, } impl WorkerPluginActivation { @@ -140,10 +143,22 @@ impl WorkerPluginActivation { /// Consumes the activation; deregistration runs from `Drop`. pub fn clear(self) {} + /// Returns the host-owned inference providers installed by this activation. + #[doc(hidden)] + pub fn inference_providers(&self) -> InferenceProviderRegistry { + self.inference_providers.clone() + } + pub(crate) fn deregister_plugin_kinds_checked(&mut self) -> DynamicPluginTeardownOutcome { deregister_tracked_registrations_checked(&mut self.plugin_registrations, "worker") } + pub(crate) fn deregister_inference_providers_checked( + &mut self, + ) -> DynamicPluginTeardownOutcome { + deregister_inference_providers_checked(&mut self.inference_provider_registrations) + } + pub(crate) fn shutdown_plugins_checked(&self) -> DynamicPluginTeardownOutcome { let mut outcome = DynamicPluginTeardownOutcome::success(); for plugin in self.plugins.iter().rev() { @@ -155,6 +170,7 @@ impl WorkerPluginActivation { impl Drop for WorkerPluginActivation { fn drop(&mut self) { + let _ = deregister_inference_providers_checked(&mut self.inference_provider_registrations); for (plugin_kind, registration_id) in self.plugin_registrations.iter().rev() { let _ = deregister_plugin_registration_checked(plugin_kind, *registration_id); } @@ -166,22 +182,43 @@ impl Drop for WorkerPluginActivation { /// The returned activation must be kept alive until after active plugin /// configuration has been cleared. pub fn load_worker_plugins(specs: I) -> crate::plugin::Result +where + I: IntoIterator, +{ + load_worker_plugins_with_inference_providers(specs, InferenceProviderRegistry::default()) +} + +/// Loads worker plugins into an existing host inference-provider registry. +#[doc(hidden)] +pub fn load_worker_plugins_with_inference_providers( + specs: I, + inference_providers: InferenceProviderRegistry, +) -> crate::plugin::Result where I: IntoIterator, { let mut activation = WorkerPluginActivation { plugins: Vec::new(), plugin_registrations: Vec::new(), + inference_providers: inference_providers.clone(), + inference_provider_registrations: Vec::new(), }; for spec in specs { let instance = load_one_worker_plugin(&spec)?; + let inference_provider_registrations = + instance.install_inference_providers(&inference_providers)?; let plugin_kind = instance.plugin_kind.clone(); + // Transfer ownership before the next fallible registration so Drop can + // unwind providers and the worker process on partial activation. + activation.plugins.push(instance.clone()); + activation + .inference_provider_registrations + .extend(inference_provider_registrations); let registration_id = register_plugin_tracked(Arc::new(WorkerPluginAdapter { plugin_kind: plugin_kind.clone(), allows_multiple_components: instance.allows_multiple_components, instance: instance.clone(), }))?; - activation.plugins.push(instance); activation .plugin_registrations .push((plugin_kind, registration_id)); @@ -1043,6 +1080,42 @@ fn clear_host_python_environment(command: &mut Command) { } impl WorkerPluginInstance { + fn install_inference_providers( + &self, + registry: &InferenceProviderRegistry, + ) -> crate::plugin::Result> { + let mut registrations = Vec::new(); + for registration in &self.registrations { + let surface = RegistrationSurface::try_from(registration.surface).map_err(|_| { + PluginError::RegistrationFailed(format!( + "worker plugin '{}' returned unsupported registration surface {}", + self.plugin_kind, registration.surface + )) + })?; + if surface != RegistrationSurface::InferenceProvider { + continue; + } + let callback_name = registration.local_name.clone(); + let provider_name = format!("{}/{}", self.plugin_kind, callback_name); + let callback = self.clone_for_callback(); + let descriptor = + InferenceProviderDescriptor::new(provider_name, registration.contract.clone())?; + match registry.register( + descriptor, + Arc::new(move |request, timeout| { + callback.invoke_inference_provider(&callback_name, request, timeout) + }), + ) { + Ok(registration) => registrations.push(registration), + Err(error) => { + let _ = deregister_inference_providers_checked(&mut registrations); + return Err(error); + } + } + } + Ok(registrations) + } + fn install_registrations( &self, ctx: &mut PluginRegistrationContext, @@ -1082,6 +1155,10 @@ impl WorkerPluginInstance { | RegistrationSurface::LlmStreamExecutionIntercept => { self.install_llm_registration(ctx, registration, surface)? } + RegistrationSurface::InferenceProvider => { + // Providers are installed during host bootstrap so components + // can resolve their contracts before runtime callbacks register. + } RegistrationSurface::Unspecified => { return Err(PluginError::RegistrationFailed(format!( "worker plugin '{}' returned unspecified registration surface", @@ -1369,6 +1446,36 @@ struct WorkerPluginCallback { } impl WorkerPluginCallback { + fn invoke_inference_provider( + &self, + registration_name: &str, + value: Json, + timeout: Duration, + ) -> crate::plugin::Result { + let request = self.base_request( + registration_name, + RegistrationSurface::InferenceProvider, + None, + Some(invoke_request_payload::Payload::Provider( + json_envelope_infallible(JSON_SCHEMA, &value), + )), + ); + let response = block_on_handle( + &self.runtime, + self.invoke_async_with_timeout(request, timeout), + ) + .map_err(|error| { + PluginError::RegistrationFailed(format!( + "inference provider '{registration_name}' invocation failed: {error}" + )) + })?; + json_from_invoke_response(response).map_err(|error| { + PluginError::RegistrationFailed(format!( + "inference provider '{registration_name}' returned an invalid response: {error}" + )) + }) + } + fn log_callback_fallback(&self, callback_name: &str, surface: RegistrationSurface) { log::warn!( target: "nemo_relay.worker", @@ -1381,6 +1488,33 @@ impl WorkerPluginCallback { } } +fn deregister_inference_providers_checked( + registrations: &mut Vec, +) -> DynamicPluginTeardownOutcome { + let mut outcome = DynamicPluginTeardownOutcome::success(); + for mut registration in std::mem::take(registrations).into_iter().rev() { + let name = registration.name().to_string(); + match registration.deregister_checked() { + Ok(PluginDeregistrationOutcome::Removed) => {} + Ok(PluginDeregistrationOutcome::Missing) => outcome.record_error( + format!("inference provider '{name}' was not registered during teardown"), + true, + ), + Ok(PluginDeregistrationOutcome::Replaced) => outcome.record_error( + format!( + "inference provider '{name}' was replaced during teardown and was left registered" + ), + true, + ), + Err(error) => outcome.record_error( + format!("failed to deregister inference provider '{name}': {error}"), + false, + ), + } + } + outcome +} + struct WorkerInvocationGuard { runtime: tokio::runtime::Handle, client: PluginWorkerClient, @@ -2948,6 +3082,19 @@ fn validate_registration_plan( "worker plugin '{plugin_id}' returned unspecified registration surface" ))); } + let contract = registration.contract.trim(); + if surface == RegistrationSurface::InferenceProvider && contract.is_empty() { + return Err(PluginError::RegistrationFailed(format!( + "worker plugin '{plugin_id}' returned inference provider '{}' without a contract", + registration.local_name + ))); + } + if surface != RegistrationSurface::InferenceProvider && !contract.is_empty() { + return Err(PluginError::RegistrationFailed(format!( + "worker plugin '{plugin_id}' returned a contract for non-provider registration '{}'", + registration.local_name + ))); + } } Ok(()) } diff --git a/crates/core/src/plugin/inference.rs b/crates/core/src/plugin/inference.rs new file mode 100644 index 000000000..fa0b6a43b --- /dev/null +++ b/crates/core/src/plugin/inference.rs @@ -0,0 +1,219 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Host-owned inference-provider services used by plugin components. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, RwLock}; +use std::time::Duration; + +use serde_json::Value as Json; + +use super::{PluginDeregistrationOutcome, PluginError, Result}; + +/// Versioned JSON request-response callback implemented by an inference provider. +#[doc(hidden)] +pub type InferenceProviderFn = Arc Result + Send + Sync + 'static>; + +/// Stable identity and request-response contract for one inference provider. +#[doc(hidden)] +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InferenceProviderDescriptor { + name: String, + contract: String, +} + +impl InferenceProviderDescriptor { + /// Creates a provider descriptor after validating its stable identifiers. + pub fn new(name: impl Into, contract: impl Into) -> Result { + let name = normalized_identifier(name.into(), "inference provider name")?; + let contract = normalized_identifier(contract.into(), "inference provider contract")?; + Ok(Self { name, contract }) + } + + /// Returns the host-qualified provider name. + pub fn name(&self) -> &str { + &self.name + } + + /// Returns the versioned request-response contract identifier. + pub fn contract(&self) -> &str { + &self.contract + } +} + +/// Resolved inference provider whose contract has already been checked. +#[doc(hidden)] +#[derive(Clone)] +pub struct InferenceProvider { + descriptor: InferenceProviderDescriptor, + callback: InferenceProviderFn, +} + +impl InferenceProvider { + /// Returns the provider descriptor. + pub fn descriptor(&self) -> &InferenceProviderDescriptor { + &self.descriptor + } + + /// Invokes the provider with the component-owned request and deadline. + pub fn invoke(&self, request: Json, timeout: Duration) -> Result { + (self.callback)(request, timeout) + } +} + +struct RegisteredInferenceProvider { + registration_id: u64, + descriptor: InferenceProviderDescriptor, + callback: InferenceProviderFn, +} + +struct InferenceProviderRegistryInner { + providers: RwLock>, + next_registration_id: AtomicU64, +} + +/// Host-scoped registry for versioned inference providers. +#[doc(hidden)] +#[derive(Clone)] +pub struct InferenceProviderRegistry { + inner: Arc, +} + +impl Default for InferenceProviderRegistry { + fn default() -> Self { + Self { + inner: Arc::new(InferenceProviderRegistryInner { + providers: RwLock::new(HashMap::new()), + next_registration_id: AtomicU64::new(1), + }), + } + } +} + +impl InferenceProviderRegistry { + /// Registers a provider and returns an ownership handle. + pub fn register( + &self, + descriptor: InferenceProviderDescriptor, + callback: InferenceProviderFn, + ) -> Result { + let mut providers = self.inner.providers.write().map_err(|error| { + PluginError::Internal(format!( + "inference provider registry lock poisoned: {error}" + )) + })?; + if providers.contains_key(descriptor.name()) { + return Err(PluginError::RegistrationFailed(format!( + "inference provider '{}' is already registered", + descriptor.name() + ))); + } + let registration_id = self + .inner + .next_registration_id + .fetch_add(1, Ordering::Relaxed); + let name = descriptor.name().to_string(); + providers.insert( + name.clone(), + RegisteredInferenceProvider { + registration_id, + descriptor, + callback, + }, + ); + Ok(InferenceProviderRegistration { + registry: self.clone(), + name, + registration_id: Some(registration_id), + }) + } + + /// Resolves a provider only when its declared contract exactly matches. + pub fn resolve(&self, name: &str, expected_contract: &str) -> Result { + let name = normalized_identifier(name.to_string(), "inference provider name")?; + let expected_contract = + normalized_identifier(expected_contract.to_string(), "inference provider contract")?; + let providers = self.inner.providers.read().map_err(|error| { + PluginError::Internal(format!( + "inference provider registry lock poisoned: {error}" + )) + })?; + let provider = providers.get(&name).ok_or_else(|| { + PluginError::NotFound(format!("inference provider '{name}' is not registered")) + })?; + if provider.descriptor.contract() != expected_contract { + return Err(PluginError::RegistrationFailed(format!( + "inference provider '{name}' implements contract '{}' but '{}' is required", + provider.descriptor.contract(), + expected_contract + ))); + } + Ok(InferenceProvider { + descriptor: provider.descriptor.clone(), + callback: Arc::clone(&provider.callback), + }) + } + + fn deregister(&self, name: &str, registration_id: u64) -> Result { + let mut providers = self.inner.providers.write().map_err(|error| { + PluginError::Internal(format!( + "inference provider registry lock poisoned: {error}" + )) + })?; + match providers.get(name) { + Some(provider) if provider.registration_id == registration_id => { + providers.remove(name); + Ok(PluginDeregistrationOutcome::Removed) + } + Some(_) => Ok(PluginDeregistrationOutcome::Replaced), + None => Ok(PluginDeregistrationOutcome::Missing), + } + } +} + +/// Owned registration for one provider in a host registry. +#[doc(hidden)] +pub struct InferenceProviderRegistration { + registry: InferenceProviderRegistry, + name: String, + registration_id: Option, +} + +impl InferenceProviderRegistration { + /// Returns the registered provider name. + pub fn name(&self) -> &str { + &self.name + } + + pub(crate) fn deregister_checked(&mut self) -> Result { + let Some(registration_id) = self.registration_id.take() else { + return Ok(PluginDeregistrationOutcome::Missing); + }; + self.registry.deregister(&self.name, registration_id) + } +} + +impl Drop for InferenceProviderRegistration { + fn drop(&mut self) { + if let Err(error) = self.deregister_checked() { + log::error!( + target: "nemo_relay.plugin", + event = "inference_provider_cleanup_failed", + provider = self.name.as_str(); + "Inference provider cleanup failed during drop: {error}" + ); + } + } +} + +fn normalized_identifier(value: String, field: &str) -> Result { + let normalized = value.trim(); + if normalized.is_empty() { + return Err(PluginError::RegistrationFailed(format!( + "{field} must not be empty" + ))); + } + Ok(normalized.to_string()) +} diff --git a/crates/core/tests/fixtures/worker_plugin/src/main.rs b/crates/core/tests/fixtures/worker_plugin/src/main.rs index a9dd41a67..99b3254ab 100644 --- a/crates/core/tests/fixtures/worker_plugin/src/main.rs +++ b/crates/core/tests/fixtures/worker_plugin/src/main.rs @@ -12,6 +12,8 @@ use serde_json::json; struct FixtureWorkerPlugin; +const DEFAULT_INFERENCE_CONTRACT: &str = "nemo.relay.pii_detection.v1"; + impl WorkerPlugin for FixtureWorkerPlugin { fn plugin_id(&self) -> &str { if std::env::var("FIXTURE_WORKER_PLUGIN_ID").as_deref() == Ok("other_worker") { @@ -57,6 +59,70 @@ impl WorkerPlugin for FixtureWorkerPlugin { ctx.register_subscriber("", |_| {}); return Ok(()); } + let inference_provider_names = config + .get("inference_provider_names") + .and_then(Json::as_array) + .map(|names| { + names + .iter() + .filter_map(Json::as_str) + .map(str::to_string) + .collect::>() + }) + .unwrap_or_else(|| { + vec![ + config + .get("inference_provider_name") + .and_then(Json::as_str) + .unwrap_or("fixture_local_model") + .to_string(), + ] + }); + for provider_name in inference_provider_names { + let callback_provider_name = provider_name.clone(); + let exit_in_local_model = fixture_flag(config, "exit_in_local_model"); + let contract = config + .get("inference_provider_contract") + .and_then(Json::as_str) + .unwrap_or(DEFAULT_INFERENCE_CONTRACT); + ctx.register_inference_provider(&provider_name, contract, move |request| { + let provider_name = callback_provider_name.clone(); + async move { + if exit_in_local_model { + std::process::exit(45); + } + if let Some(delay_ms) = request.get("delay_ms").and_then(Json::as_u64) { + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + } + if let Some(texts) = request.get("texts").and_then(Json::as_array) { + let detections = texts + .iter() + .filter_map(|item| { + let text_id = item.get("id")?.as_u64()?; + let text = item.get("text")?.as_str()?; + let start_utf8 = text.find("PRIVATE")?; + Some(json!({ + "text_id": text_id, + "start_utf8": start_utf8, + "end_utf8": start_utf8 + "PRIVATE".len(), + "label": "fixture_private", + "score": 1.0 + })) + }) + .collect::>(); + return Ok(json!({"version": 1, "detections": detections})); + } + Ok(json!({ + "version": 1, + "request": request, + "provider": provider_name + })) + } + }); + } + if fixture_flag(config, "provider_only") { + return Ok(()); + } let runtime = ctx .runtime() diff --git a/crates/core/tests/integration/worker_plugin_tests.rs b/crates/core/tests/integration/worker_plugin_tests.rs index 4387b048a..b83940da2 100644 --- a/crates/core/tests/integration/worker_plugin_tests.rs +++ b/crates/core/tests/integration/worker_plugin_tests.rs @@ -24,17 +24,20 @@ use nemo_relay::codec::traits::LlmCodec; use nemo_relay::error::Result as FlowResult; use nemo_relay::plugin::dynamic::{ DynamicPluginActivationSpec, DynamicPluginKind, PluginHostActivation, WorkerPluginActivation, - WorkerPluginLoadSpec, load_worker_plugins, + WorkerPluginLoadSpec, load_worker_plugins, load_worker_plugins_with_inference_providers, }; use nemo_relay::plugin::{ - PluginComponentSpec, PluginConfig, clear_plugin_configuration, initialize_plugins_exact, - list_plugin_kinds, + InferenceProviderDescriptor, InferenceProviderRegistry, PluginComponentSpec, PluginConfig, + clear_plugin_configuration, initialize_plugins_exact, + initialize_plugins_exact_with_inference_providers, list_plugin_kinds, }; use serde_json::{Map, Value as Json, json}; use sha2::{Digest, Sha256}; use tempfile::TempDir; use uuid::Uuid; +const PII_DETECTION_CONTRACT: &str = "nemo.relay.pii_detection.v1"; +const EXAMPLE_ECHO_CONTRACT: &str = "examples.python_grpc_worker.echo.v1"; static WORKER_PLUGIN_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); fn enable_operational_logs() { @@ -50,6 +53,203 @@ fn worker_activation_with_no_specs_is_empty() { activation.clear(); } +#[tokio::test(flavor = "multi_thread")] +async fn worker_inference_provider_is_preinstalled_times_out_and_clears() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + let activation = load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + environment_ref: None, + config: Map::new(), + }]) + .expect("worker plugin should load"); + let registry = activation.inference_providers(); + + // Providers must be available before static consumers initialize. + let provider = registry + .resolve("fixture_worker/fixture_local_model", PII_DETECTION_CONTRACT) + .expect("worker provider should be installed"); + assert_eq!( + provider + .invoke( + json!({"text": "private"}), + std::time::Duration::from_secs(1) + ) + .expect("worker provider should return JSON"), + json!({ + "version": 1, + "request": {"text": "private"}, + "provider": "fixture_local_model" + }) + ); + let timeout = provider + .invoke( + json!({"delay_ms": 100}), + std::time::Duration::from_millis(5), + ) + .expect_err("worker provider should honor the caller deadline") + .to_string(); + assert!(timeout.contains("timed out"), "{timeout}"); + + activation.clear(); + assert!( + registry + .resolve("fixture_worker/fixture_local_model", PII_DETECTION_CONTRACT) + .is_err(), + "provider should be removed when the worker activation clears" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn worker_clear_fails_an_in_flight_local_model_call_without_hanging() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + let activation = load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + environment_ref: None, + config: Map::new(), + }]) + .expect("worker plugin should load"); + let registry = activation.inference_providers(); + let provider = registry + .resolve("fixture_worker/fixture_local_model", PII_DETECTION_CONTRACT) + .expect("worker provider should be installed"); + let invocation = std::thread::spawn(move || { + provider.invoke( + json!({"delay_ms": 5_000}), + std::time::Duration::from_secs(10), + ) + }); + + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + activation.clear(); + + let error = invocation + .join() + .expect("provider invocation thread should join") + .expect_err("clearing the worker must fail its in-flight call") + .to_string(); + assert!( + error.contains("invocation failed") || error.contains("cancel"), + "{error}" + ); + assert!( + registry + .resolve("fixture_worker/fixture_local_model", PII_DETECTION_CONTRACT) + .is_err(), + "provider should remain deregistered after concurrent clear" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn worker_provider_rolls_back_after_later_plugin_registration_failure() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + let first_provider = "fixture_local_model_first"; + let first_provider_key = format!("fixture_worker/{first_provider}"); + let registry = InferenceProviderRegistry::default(); + let first = load_worker_plugins_with_inference_providers( + [WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + environment_ref: None, + config: Map::from_iter([("inference_provider_name".into(), json!(first_provider))]), + }], + registry.clone(), + ) + .expect("first worker plugin should load"); + + let second_provider = "fixture_local_model_rollback"; + let second_provider_key = format!("fixture_worker/{second_provider}"); + let second = load_worker_plugins_with_inference_providers( + [WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + environment_ref: None, + config: Map::from_iter([("inference_provider_name".into(), json!(second_provider))]), + }], + registry.clone(), + ); + assert!( + second.is_err(), + "duplicate plugin kind should fail after the second provider is installed" + ); + assert!( + registry + .resolve(&second_provider_key, PII_DETECTION_CONTRACT) + .is_err(), + "the second provider must be rolled back with its failed activation" + ); + assert!( + registry + .resolve(&first_provider_key, PII_DETECTION_CONTRACT) + .is_ok(), + "rollback must not remove the first activation's provider" + ); + + first.clear(); + assert!( + registry + .resolve(&first_provider_key, PII_DETECTION_CONTRACT) + .is_err() + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn worker_provider_rolls_back_earlier_provider_after_same_worker_conflict() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + let first_provider_key = "fixture_worker/fixture_local_model_unique"; + let conflicting_provider_key = "fixture_worker/fixture_local_model_conflict"; + let registry = InferenceProviderRegistry::default(); + let _existing_registration = registry + .register( + InferenceProviderDescriptor::new(conflicting_provider_key, PII_DETECTION_CONTRACT) + .unwrap(), + Arc::new(|request, _| Ok(json!({"existing": request}))), + ) + .expect("conflicting provider fixture should register"); + + let activation = load_worker_plugins_with_inference_providers( + [WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + environment_ref: None, + config: Map::from_iter([( + "inference_provider_names".into(), + json!(["fixture_local_model_unique", "fixture_local_model_conflict"]), + )]), + }], + registry.clone(), + ); + + assert!( + activation.is_err(), + "the worker activation should fail on its second provider" + ); + assert!( + registry + .resolve(first_provider_key, PII_DETECTION_CONTRACT) + .is_err(), + "an earlier provider from the failed worker must be rolled back" + ); + let existing = registry + .resolve(conflicting_provider_key, PII_DETECTION_CONTRACT) + .expect("the existing conflicting provider must remain registered"); + assert_eq!( + existing + .invoke(json!({"value": 1}), std::time::Duration::from_secs(1)) + .expect("existing provider should remain callable"), + json!({"existing": {"value": 1}}) + ); +} + #[tokio::test] async fn plugin_host_activation_owns_worker_lifecycle() { let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; @@ -1097,6 +1297,7 @@ async fn python_worker_host_runtime_mark_and_mutated_request_round_trip() { config: config.clone(), }]) .expect("managed Python worker should load"); + let inference_providers = activation.inference_providers(); let mut cleanup = PythonWorkerCleanup::new(activation); let mut plugin_config = PluginConfig::default(); @@ -1105,7 +1306,7 @@ async fn python_worker_host_runtime_mark_and_mutated_request_round_trip() { enabled: true, config, }); - initialize_plugins_exact(plugin_config) + initialize_plugins_exact_with_inference_providers(plugin_config, inference_providers.clone()) .await .expect("managed Python worker should initialize"); @@ -1125,6 +1326,24 @@ async fn python_worker_host_runtime_mark_and_mutated_request_round_trip() { rewritten["_nemo_relay_plugin"]["tag"], "managed-environment" ); + let inference_provider = inference_providers + .resolve("examples.python_grpc_worker/echo", EXAMPLE_ECHO_CONTRACT) + .expect("Python worker should expose its inference provider"); + assert_eq!( + inference_provider + .invoke( + json!({"version": 1, "texts": [{"id": 0, "text": "private"}]}), + std::time::Duration::from_secs(1), + ) + .expect("Python inference provider should round-trip JSON"), + json!({ + "provider": "python_grpc_worker", + "request": { + "version": 1, + "texts": [{"id": 0, "text": "private"}], + }, + }), + ); flush_subscribers().expect("Python callback mark should flush"); find_event( &events.lock().unwrap(), @@ -1236,6 +1455,7 @@ async fn load_and_initialize_fixture(config: Map) -> LoadedWorker config: config.clone(), }]) .expect("worker plugin should load"); + let inference_providers = activation.inference_providers(); let mut plugin_config = PluginConfig::default(); plugin_config.components.push(PluginComponentSpec { @@ -1243,7 +1463,7 @@ async fn load_and_initialize_fixture(config: Map) -> LoadedWorker enabled: true, config, }); - initialize_plugins_exact(plugin_config) + initialize_plugins_exact_with_inference_providers(plugin_config, inference_providers) .await .expect("worker plugin should initialize"); diff --git a/crates/core/tests/unit/dynamic_worker_tests.rs b/crates/core/tests/unit/dynamic_worker_tests.rs index 7754238dc..0d250926d 100644 --- a/crates/core/tests/unit/dynamic_worker_tests.rs +++ b/crates/core/tests/unit/dynamic_worker_tests.rs @@ -246,6 +246,7 @@ fn registration_plan_and_scope_type_helpers_validate_edges() { surface: RegistrationSurface::Subscriber as i32, priority: 0, break_chain: false, + contract: String::new(), }], error: None, }, @@ -261,6 +262,7 @@ fn registration_plan_and_scope_type_helpers_validate_edges() { surface: 999, priority: 0, break_chain: false, + contract: String::new(), }], error: None, }, @@ -280,6 +282,7 @@ fn registration_plan_and_scope_type_helpers_validate_edges() { surface: RegistrationSurface::Unspecified as i32, priority: 0, break_chain: false, + contract: String::new(), }], error: None, }, @@ -291,6 +294,44 @@ fn registration_plan_and_scope_type_helpers_validate_edges() { .contains("unspecified registration surface") ); + let missing_contract = validate_registration_plan( + "fixture_worker", + &RegisterResponse { + registrations: vec![registration( + RegistrationSurface::InferenceProvider, + "detector", + )], + error: None, + }, + ) + .expect_err("inference providers must declare a contract"); + assert!(missing_contract.to_string().contains("without a contract")); + + let contract_on_middleware = validate_registration_plan( + "fixture_worker", + &RegisterResponse { + registrations: vec![Registration { + contract: "test.detector.v1".into(), + ..registration(RegistrationSurface::Subscriber, "subscriber") + }], + error: None, + }, + ) + .expect_err("middleware registrations must not declare provider contracts"); + assert!(contract_on_middleware.to_string().contains("non-provider")); + + validate_registration_plan( + "fixture_worker", + &RegisterResponse { + registrations: vec![Registration { + contract: "test.detector.v1".into(), + ..registration(RegistrationSurface::InferenceProvider, "detector") + }], + error: None, + }, + ) + .expect("versioned inference provider contract should be accepted"); + let cases = [ (ProtoScopeType::Agent, crate::api::scope::ScopeType::Agent), ( @@ -2148,6 +2189,7 @@ fn registration(surface: RegistrationSurface, local_name: &str) -> Registration surface: surface as i32, priority: 0, break_chain: false, + contract: String::new(), } } diff --git a/crates/core/tests/unit/inference_provider_tests.rs b/crates/core/tests/unit/inference_provider_tests.rs new file mode 100644 index 000000000..5e9ccc5e7 --- /dev/null +++ b/crates/core/tests/unit/inference_provider_tests.rs @@ -0,0 +1,117 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; +use std::time::Duration; + +use serde_json::json; + +use crate::plugin::{InferenceProviderDescriptor, InferenceProviderRegistry}; + +#[test] +fn provider_round_trips_json_and_receives_deadline() { + let registry = InferenceProviderRegistry::default(); + let _registration = registry + .register( + InferenceProviderDescriptor::new("test-provider", "test.echo.v1").unwrap(), + Arc::new(|request, timeout| { + assert_eq!(timeout, Duration::from_millis(25)); + Ok(json!({"request": request})) + }), + ) + .unwrap(); + + let provider = registry.resolve("test-provider", "test.echo.v1").unwrap(); + assert_eq!( + provider + .invoke(json!({"text": "hello"}), Duration::from_millis(25)) + .unwrap(), + json!({"request": {"text": "hello"}}) + ); +} + +#[test] +fn registration_owns_provider_lifetime() { + let registry = InferenceProviderRegistry::default(); + let registration = registry + .register( + InferenceProviderDescriptor::new("owned-provider", "test.echo.v1").unwrap(), + Arc::new(|request, _| Ok(request)), + ) + .unwrap(); + + assert!(registry.resolve("owned-provider", "test.echo.v1").is_ok()); + drop(registration); + assert!(registry.resolve("owned-provider", "test.echo.v1").is_err()); +} + +#[test] +fn duplicate_provider_names_are_rejected() { + let registry = InferenceProviderRegistry::default(); + let _registration = registry + .register( + InferenceProviderDescriptor::new("duplicate-provider", "test.echo.v1").unwrap(), + Arc::new(|request, _| Ok(request)), + ) + .unwrap(); + let duplicate = registry + .register( + InferenceProviderDescriptor::new("duplicate-provider", "test.other.v1").unwrap(), + Arc::new(|request, _| Ok(request)), + ) + .err() + .expect("duplicate provider names must fail"); + + assert!(duplicate.to_string().contains("already registered")); +} + +#[test] +fn provider_names_are_normalized_consistently() { + let registry = InferenceProviderRegistry::default(); + let _registration = registry + .register( + InferenceProviderDescriptor::new(" normalized-provider ", " test.echo.v1 ") + .unwrap(), + Arc::new(|request, _| Ok(request)), + ) + .unwrap(); + + assert!( + registry + .resolve(" normalized-provider ", " test.echo.v1 ") + .is_ok() + ); +} + +#[test] +fn provider_contract_mismatch_is_rejected_before_invocation() { + let registry = InferenceProviderRegistry::default(); + let _registration = registry + .register( + InferenceProviderDescriptor::new("detector", "test.detector.v1").unwrap(), + Arc::new(|request, _| Ok(request)), + ) + .unwrap(); + + let error = registry + .resolve("detector", "test.embedding.v1") + .err() + .expect("mismatched contracts must fail"); + assert!(error.to_string().contains("test.detector.v1")); + assert!(error.to_string().contains("test.embedding.v1")); +} + +#[test] +fn registries_isolate_provider_names_between_hosts() { + let first = InferenceProviderRegistry::default(); + let second = InferenceProviderRegistry::default(); + let _first_registration = first + .register( + InferenceProviderDescriptor::new("shared-name", "test.echo.v1").unwrap(), + Arc::new(|request, _| Ok(request)), + ) + .unwrap(); + + assert!(first.resolve("shared-name", "test.echo.v1").is_ok()); + assert!(second.resolve("shared-name", "test.echo.v1").is_err()); +} diff --git a/crates/core/tests/unit/plugin_tests.rs b/crates/core/tests/unit/plugin_tests.rs index b39abb9d2..2cab29057 100644 --- a/crates/core/tests/unit/plugin_tests.rs +++ b/crates/core/tests/unit/plugin_tests.rs @@ -26,6 +26,7 @@ struct RecordingPlugin; struct ReplacementPlugin; struct RestoreFailPlugin; struct RestoreBreakPlugin; +struct ProviderAwarePlugin; struct PartialFailPlugin; struct VanishingPlugin; struct BlockingPlugin { @@ -55,11 +56,16 @@ static PARTIAL_FAIL_ROLLBACKS: AtomicUsize = AtomicUsize::new(0); static RESTORE_FAIL_REGISTRATIONS: AtomicUsize = AtomicUsize::new(0); static RESTORE_BREAK_REGISTRATIONS: AtomicUsize = AtomicUsize::new(0); static REPLACEMENT_REGISTRATIONS: AtomicUsize = AtomicUsize::new(0); +static INFERENCE_PROVIDER_VALUES: OnceLock>> = OnceLock::new(); fn recorded_names() -> &'static Mutex> { RECORDED_NAMES.get_or_init(|| Mutex::new(Vec::new())) } +fn inference_provider_values() -> &'static Mutex> { + INFERENCE_PROVIDER_VALUES.get_or_init(|| Mutex::new(Vec::new())) +} + fn lock_runtime_owner() -> std::sync::MutexGuard<'static, ()> { crate::shared_runtime::runtime_owner_test_mutex() .lock() @@ -305,6 +311,40 @@ impl Plugin for RestoreBreakPlugin { } } +impl Plugin for ProviderAwarePlugin { + fn plugin_kind(&self) -> &str { + "provider-aware.plugin" + } + + fn validate(&self, _plugin_config: &Map) -> Vec { + vec![] + } + + fn register<'a>( + &'a self, + _plugin_config: &Map, + ctx: &'a mut PluginRegistrationContext, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + let provider = ctx.inference_provider("shared-provider", "test.echo.v1")?; + let response = provider.invoke(json!({}), std::time::Duration::from_secs(1))?; + let source = response + .get("source") + .and_then(Json::as_str) + .ok_or_else(|| { + PluginError::RegistrationFailed( + "provider-aware.plugin received an invalid response".into(), + ) + })?; + inference_provider_values() + .lock() + .unwrap() + .push(source.to_string()); + Ok(()) + }) + } +} + impl Plugin for PartialFailPlugin { fn plugin_kind(&self) -> &str { "partial.fail.plugin" @@ -477,12 +517,14 @@ fn reset_global() { RESTORE_FAIL_REGISTRATIONS.store(0, Ordering::SeqCst); RESTORE_BREAK_REGISTRATIONS.store(0, Ordering::SeqCst); REPLACEMENT_REGISTRATIONS.store(0, Ordering::SeqCst); + inference_provider_values().lock().unwrap().clear(); let _ = deregister_plugin("test.plugin"); let _ = deregister_plugin("singleton.plugin"); let _ = deregister_plugin("recording.plugin"); let _ = deregister_plugin("replacement.plugin"); let _ = deregister_plugin("restore.fail.plugin"); let _ = deregister_plugin("restore.break.plugin"); + let _ = deregister_plugin("provider-aware.plugin"); let _ = deregister_plugin("partial.fail.plugin"); let _ = deregister_plugin("vanishing.plugin"); let _ = deregister_plugin("blocking.plugin"); @@ -1110,6 +1152,60 @@ fn test_initialize_plugins_restores_previous_configuration_after_failed_replacem reset_global(); } +#[test] +fn test_failed_replacement_restores_previous_inference_provider_registry() { + let _guard = lock_runtime_owner(); + reset_global(); + register_plugin(Arc::new(ProviderAwarePlugin)).unwrap(); + register_plugin(Arc::new(RestoreFailPlugin)).unwrap(); + + let previous_registry = InferenceProviderRegistry::default(); + let _previous_provider = previous_registry + .register( + InferenceProviderDescriptor::new("shared-provider", "test.echo.v1").unwrap(), + Arc::new(|_, _| Ok(json!({"source": "previous"}))), + ) + .unwrap(); + let replacement_registry = InferenceProviderRegistry::default(); + let _replacement_provider = replacement_registry + .register( + InferenceProviderDescriptor::new("shared-provider", "test.echo.v1").unwrap(), + Arc::new(|_, _| Ok(json!({"source": "replacement"}))), + ) + .unwrap(); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime + .block_on(initialize_plugins_exact_with_inference_providers( + PluginConfig { + components: vec![PluginComponentSpec::new("provider-aware.plugin")], + ..PluginConfig::default() + }, + previous_registry, + )) + .unwrap(); + + let error = runtime + .block_on(initialize_plugins_exact_with_inference_providers( + PluginConfig { + components: vec![PluginComponentSpec::new("restore.fail.plugin")], + ..PluginConfig::default() + }, + replacement_registry, + )) + .unwrap_err(); + assert!(error.to_string().contains("refused to initialize")); + assert_eq!( + *inference_provider_values().lock().unwrap(), + vec!["previous", "previous"] + ); + + reset_global(); +} + #[test] fn test_initialize_plugins_restores_previous_configuration_after_replacement_panic() { let _guard = lock_runtime_owner(); @@ -1318,6 +1414,7 @@ fn test_checked_teardown_reports_unremoved_registrations() { )) }), )], + InferenceProviderRegistry::default(), ) .unwrap(); @@ -1343,6 +1440,7 @@ fn test_legacy_clear_retains_mutation_owner_after_incomplete_teardown() { "stale-callback", Box::new(|| panic!("fixture deregistration panicked")), )], + InferenceProviderRegistry::default(), ) .unwrap(); diff --git a/crates/node/pii_redaction.d.ts b/crates/node/pii_redaction.d.ts index 5a5f20e76..2df3407ea 100644 --- a/crates/node/pii_redaction.d.ts +++ b/crates/node/pii_redaction.d.ts @@ -26,10 +26,23 @@ export interface LocalModelConfig { backend?: string; model_id?: string; detector_profile?: string; + target_paths?: string[]; + target_path_patterns?: string[]; + min_score?: number; + excluded_labels?: string[]; + replacement?: string; allow_network?: boolean; max_latency_ms?: number; } +export interface ProfileConfig { + enabled?: boolean; + mode?: 'builtin' | 'local_model' | string; + priority?: number; + builtin?: BuiltinConfig; + local?: LocalModelConfig; +} + export interface Config { version?: number; mode?: 'builtin' | 'local_model' | string; @@ -40,6 +53,7 @@ export interface Config { mark?: boolean; priority?: number; codec?: 'openai_chat' | 'openai_responses' | 'anthropic_messages' | string; + profiles?: ProfileConfig[]; builtin?: BuiltinConfig; local?: LocalModelConfig; policy?: ConfigPolicy; @@ -57,8 +71,10 @@ export declare const PII_REDACTION_PLUGIN_KIND: 'pii_redaction'; export declare function defaultConfig(): Config; /** Create deterministic built-in redaction backend settings with defaults applied. */ export declare function builtinConfig(config?: BuiltinConfig): BuiltinConfig; -/** Create future local-model backend settings with defaults applied. */ +/** Create worker-backed local-model provider settings. */ export declare function localModelConfig(config?: LocalModelConfig): LocalModelConfig; +/** Create one ordered redaction profile with defaults applied. */ +export declare function profileConfig(config?: ProfileConfig): ProfileConfig; /** Wrap PII redaction config as a top-level plugin component. */ export declare function ComponentSpec( config: Config, diff --git a/crates/node/pii_redaction.js b/crates/node/pii_redaction.js index c5fc6d313..dea68d566 100644 --- a/crates/node/pii_redaction.js +++ b/crates/node/pii_redaction.js @@ -39,7 +39,7 @@ function builtinConfig(config = {}) { } /** - * Create future local-model backend settings with defaults applied. + * Create worker-backed local-model redaction settings. * * @param {object} [config={}] - Partial local-model settings to override. * @returns {object} A normalized local-model backend config object. @@ -50,6 +50,21 @@ function localModelConfig(config = {}) { }; } +/** + * Create one ordered redaction profile with defaults applied. + * + * @param {object} [config={}] - Partial profile settings to override. + * @returns {object} A normalized redaction profile object. + */ +function profileConfig(config = {}) { + return { + enabled: true, + mode: 'builtin', + priority: 100, + ...config, + }; +} + /** * Wrap PII redaction config as a top-level plugin component. * @@ -68,5 +83,6 @@ module.exports = { defaultConfig, builtinConfig, localModelConfig, + profileConfig, ComponentSpec, }; diff --git a/crates/node/tests/pii_redaction_tests.mjs b/crates/node/tests/pii_redaction_tests.mjs index 0c0edc3f7..e696e4880 100644 --- a/crates/node/tests/pii_redaction_tests.mjs +++ b/crates/node/tests/pii_redaction_tests.mjs @@ -23,6 +23,37 @@ describe('pii_redaction plugin helpers', () => { }); assert.deepEqual(piiRedaction.builtinConfig(), { action: 'remove' }); assert.deepEqual(piiRedaction.localModelConfig(), {}); + assert.deepEqual( + piiRedaction.localModelConfig({ + backend: 'acme.pii/detector', + model_id: 'pii-model-v1', + detector_profile: 'default', + target_paths: ['/message'], + target_path_patterns: ['/messages/*/content'], + min_score: 0.6, + excluded_labels: ['CITY'], + replacement: '[PRIVATE]', + allow_network: false, + max_latency_ms: 250, + }), + { + backend: 'acme.pii/detector', + model_id: 'pii-model-v1', + detector_profile: 'default', + target_paths: ['/message'], + target_path_patterns: ['/messages/*/content'], + min_score: 0.6, + excluded_labels: ['CITY'], + replacement: '[PRIVATE]', + allow_network: false, + max_latency_ms: 250, + }, + ); + assert.deepEqual(piiRedaction.profileConfig(), { + enabled: true, + mode: 'builtin', + priority: 100, + }); const component = piiRedaction.ComponentSpec({ ...piiRedaction.defaultConfig(), @@ -32,6 +63,29 @@ describe('pii_redaction plugin helpers', () => { assert.equal(component.enabled, true); }); + it('builds profile composition without legacy top-level fields', () => { + const config = { + version: 1, + codec: 'openai_chat', + profiles: [ + piiRedaction.profileConfig({ + builtin: piiRedaction.builtinConfig({ detector: 'email' }), + }), + piiRedaction.profileConfig({ + mode: 'local_model', + priority: 110, + local: piiRedaction.localModelConfig({ + backend: 'acme.pii/detector', + target_paths: ['/message'], + }), + }), + ], + }; + + assert.equal(config.mode, undefined); + assert.deepEqual(plugin.validate({ version: 1, components: [piiRedaction.ComponentSpec(config)] }).diagnostics, []); + }); + it('lists builtin pii_redaction kind and validates bad values', () => { assert.equal(plugin.listKinds().includes(piiRedaction.PII_REDACTION_PLUGIN_KIND), true); const report = plugin.validate({ diff --git a/crates/pii-redaction/Cargo.toml b/crates/pii-redaction/Cargo.toml index 006a45b59..fb2163f23 100644 --- a/crates/pii-redaction/Cargo.toml +++ b/crates/pii-redaction/Cargo.toml @@ -27,7 +27,7 @@ sha2 = "0.11" schemars = { version = "0.8", optional = true } [dev-dependencies] -nemo-relay = { workspace = true, features = ["openinference", "otel"] } +nemo-relay = { workspace = true, features = ["openinference", "otel", "worker-grpc"] } futures = "0.3" tokio = { version = "1", features = ["rt", "macros", "sync", "test-util", "rt-multi-thread", "time"] } opentelemetry_sdk = { workspace = true, features = ["trace", "testing"] } diff --git a/crates/pii-redaction/README.md b/crates/pii-redaction/README.md index 4c8528582..10f3d4cf6 100644 --- a/crates/pii-redaction/README.md +++ b/crates/pii-redaction/README.md @@ -8,7 +8,7 @@ SPDX-License-Identifier: Apache-2.0 `nemo-relay-pii-redaction` is the first-party NeMo Relay plugin crate for deterministic privacy redaction on tool and LLM observability payloads. It ships the `pii_redaction` plugin contract, a production-ready `builtin` -backend, and the future `local_model` seam for model-backed detection and +backend, and a worker-backed `local_model` seam for model-backed detection and redaction. The plugin is designed for the common case where teams want a supported, @@ -36,8 +36,8 @@ NeMo Relay PII Redaction allows you to: `openai_responses`, and `anthropic_messages`. - Remove conversational trajectory content while preserving event structure, tool-call identity, model attribution, routing, usage, and cost analytics. -- Use the `local_model` config contract and provider registration surface for - future model-backed implementations. +- Use an isolated `grpc-v1` worker as the detector behind the `local_model` + config contract without loading model dependencies into the Relay host. ## Plugin Versus Raw Middleware @@ -188,13 +188,153 @@ high-risk secrets, prefer `redact` over partial `mask` behavior. ## Local Model Seam -`local_model` is included in the plugin contract now, but no runtime -implementation ships in this crate yet. +`local_model` delegates bounded detector inference to a manifest-backed +`grpc-v1` worker. The PII component remains responsible for choosing +observability fields, decoding provider payloads, batching text, enforcing the +deadline and failure policy, validating detections, and replacing accepted +spans. -The seam exists so a future local detector/redactor backend can be added -without redesigning the public plugin surface. If `mode = "local_model"` is -configured today, the runtime expects a registered local backend provider and -fails fast if one is not installed. +Configure the provider by its host-qualified name: + +```toml +[[components]] +kind = "pii_redaction" +enabled = true + +[components.config] +mode = "local_model" +codec = "openai_chat" + +[components.config.local] +backend = "acme.pii_worker/detector" +model_id = "acme-pii-v1" +detector_profile = "default" +min_score = 0.4 +target_path_patterns = [ + "/messages/*/content", + "/messages/*/content/*/text", + "/message", + "/message/*/text", +] +replacement = "[REDACTED]" +allow_network = false +max_latency_ms = 250 +``` + +The backend name is `/`. For example, a worker +with plugin ID `acme.pii_worker` that calls +`register_inference_provider("detector", "nemo.relay.pii_detection.v1", ...)` +is selected as `acme.pii_worker/detector`. Relay verifies the PII contract, +installs worker providers before static components initialize, and removes PII +sanitizers before stopping their worker. + +Use profiles to compose deterministic and contextual detection. The lower +priority runs first: + +```toml +[components.config] +codec = "openai_chat" + +[[components.config.profiles]] +mode = "builtin" +priority = 80 + +[components.config.profiles.builtin] +action = "redact" +detector = "email" + +[[components.config.profiles]] +mode = "local_model" +priority = 90 + +[components.config.profiles.local] +backend = "nemo_relay.pii_rampart/detector" +min_score = 0.4 +max_latency_ms = 5000 +target_path_patterns = [ + "/messages/*/content", + "/messages/*/content/*/text", + "/message", + "/message/*/text", +] +``` + +`target_paths` contains exact JSON pointers. `target_path_patterns` also accepts +`*` as one complete path segment, which is useful for message arrays. When a +codec is configured, paths address the normalized request or response shape; +the content-only patterns above cover the `openai_chat` request and response +shape. Without a codec, they address the original JSON payload. When both lists +are empty, Relay inspects every string leaf and reports a configuration warning. + +Use content-only paths for contextual classifiers. Do not send model names, +tool identifiers, trace IDs, routing fields, or arbitrary provider metadata to +a classifier unless that is an explicit policy choice. + +Relay accepts detections whose confidence is at least `min_score`, which +defaults to `0.4`. `excluded_labels` is an exact, case-sensitive denylist for +provider labels that should remain visible. The host applies both settings +after validating the complete provider response; workers do not own the final +redaction policy. + +Provider failures, timeouts, malformed responses, invalid UTF-8 boundaries, +overlapping spans, and input-limit violations fail closed for the affected +batch. If a configured codec cannot decode or safely re-encode an LLM payload, +Relay omits that request or response payload from the emitted event; it does not +retry normalized selectors against the raw provider shape. +`allow_network = true` is rejected; this lane is for same-machine inference. +This setting is a configuration invariant, not a network sandbox: Relay's +worker launcher does not currently prevent a worker process from opening +sockets. Only install providers whose packaging and runtime behavior satisfy +that policy. The default deadline is 250 ms for the complete selected payload, +including every provider batch. Configuration above 60 seconds is rejected. + +### Provider Contract + +The worker receives a versioned JSON request: + +```json +{ + "version": 1, + "model_id": "acme-pii-v1", + "detector_profile": "default", + "texts": [ + {"id": 0, "text": "Contact Alice Rivera"} + ] +} +``` + +It returns detections using UTF-8 byte offsets: + +```json +{ + "version": 1, + "detections": [ + { + "text_id": 0, + "start_utf8": 8, + "end_utf8": 20, + "label": "person", + "score": 0.99 + } + ] +} +``` + +The provider performs inference only. It must not choose Relay surfaces, +traverse arbitrary event fields, or apply replacements itself. Rust and Python +workers have SDK helpers for this registration. Other languages can implement +the same `grpc-v1` protobuf contract directly; Rust, Python, and Node hosts all +consume it through the shared core runtime. + +### Optional Rampart Provider + +The source tree includes an optional +[Rampart worker](./providers/rampart/README.md) that implements this provider +contract with a pinned ONNX token-classification model. It runs in a +Relay-managed Python worker process, keeps ONNX dependencies out of the host, +and complements the built-in deterministic recognizers. The model is prefetched +separately, then resolved and integrity-verified from the local cache at +activation. It is not distributed in the Relay package. ## Documentation diff --git a/crates/pii-redaction/providers/rampart/MANIFEST.in b/crates/pii-redaction/providers/rampart/MANIFEST.in new file mode 100644 index 000000000..af375315e --- /dev/null +++ b/crates/pii-redaction/providers/rampart/MANIFEST.in @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +include config.schema.json +include relay-plugin.toml diff --git a/crates/pii-redaction/providers/rampart/README.md b/crates/pii-redaction/providers/rampart/README.md new file mode 100644 index 000000000..91d49d4b6 --- /dev/null +++ b/crates/pii-redaction/providers/rampart/README.md @@ -0,0 +1,161 @@ + + +# Rampart PII Provider + +This optional manifest-backed Python worker runs the +[`nationaldesignstudio/rampart`](https://huggingface.co/nationaldesignstudio/rampart) +ONNX token classifier behind NeMo Relay's `pii_redaction.local_model` backend. +Relay owns field selection, event sanitization, replacement, deadlines, and +fail-closed behavior. The worker performs detector inference only. + +The worker runs as a local child process over Relay's `grpc-v1` protocol. Its +Python, ONNX Runtime, NumPy, tokenizer, and model-cache dependencies remain in a +Relay-managed virtual environment rather than the Relay host process. Process +isolation is not a security sandbox. + +## Install + +From this directory: + +```bash +uvx --from . nemo-relay-pii-rampart-prefetch +nemo-relay plugins add ./relay-plugin.toml +nemo-relay plugins enable nemo_relay.pii_rampart +``` + +If the provider package is already installed, run +`nemo-relay-pii-rampart-prefetch` directly. `plugins add` creates a separate +Relay-managed Python environment from the same source directory. + +Rampart activation is offline-only. It requires the pinned 14.7 MB model +snapshot to already exist in the Hugging Face cache. Model acquisition is an +explicit setup step and never occurs during activation or inference. + +`model_id` normally remains `nationaldesignstudio/rampart`. The worker rejects +other repository identifiers and revisions because its integrity manifest pins +one supported snapshot. To load that same snapshot from a local directory, set +`model_id` to an absolute path or use explicit `./`, `../`, or `~/` syntax so a +relative directory cannot shadow the repository identifier. The PII +component's optional `local.model_id` remains the logical +`nationaldesignstudio/rampart` identifier even when worker storage uses a local +path. + +On hosts with slow or restricted access to Hugging Face, populate the shared +cache before enabling the plugin: + +```bash +nemo-relay-pii-rampart-prefetch +``` + +Run that command as the same operating-system user that runs Relay. If +`cache_dir` is configured for the plugin, pass the same value with +`--cache-dir`. The prefetch command and activation both verify SHA-256 digests +for every runtime model and tokenizer file. A missing or modified file blocks +activation. + +## Configure + +Add worker settings to the `[[plugins.dynamic]]` record created by `plugins +add`: + +```toml +[plugins.dynamic.config] +local_files_only = true +max_windows_per_request = 128 +inference_batch_size = 16 +max_pending_requests = 8 +``` + +`inference_batch_size` is an upper bound. The worker batches short token +windows together but automatically reduces the batch width for longer windows +to bound ONNX intermediate memory. + +Add the PII component to the same `plugins.toml`: + +```toml +[[components]] +kind = "pii_redaction" +enabled = true + +[components.config] +codec = "openai_chat" + +[[components.config.profiles]] +mode = "builtin" +priority = 70 + +[components.config.profiles.builtin] +action = "redact" +detector = "email" + +[[components.config.profiles]] +mode = "builtin" +priority = 80 + +[components.config.profiles.builtin] +action = "redact" +detector = "credit_card" + +[[components.config.profiles]] +mode = "local_model" +priority = 90 + +[components.config.profiles.local] +backend = "nemo_relay.pii_rampart/detector" +model_id = "nationaldesignstudio/rampart" +detector_profile = "default" +allow_network = false +max_latency_ms = 5000 +min_score = 0.4 +replacement = "[REDACTED]" +target_path_patterns = [ + "/messages/*/content", + "/messages/*/content/*/text", + "/message", + "/message/*/text", +] +``` + +`allow_network = false` means provider inference is local. It does not sandbox +the worker. `local_files_only` must remain `true`; activation-time model +acquisition is not supported. + +Rampart is the contextual detector lane, not a replacement for deterministic +recognizers. Configure built-in PII profiles for structured values and the +local-model profile for names and contextual identifiers. Keep the local-model +profile limited to normalized content paths. Classifying every string leaf can +produce false positives on model names, region names, UUIDs, trace IDs, and +other machine identifiers. Relay, not the worker, applies `min_score` and +optional `excluded_labels` policy after validating the provider response. + +## Runtime Bounds + +- At most 64 texts and 64 KiB of UTF-8 text are accepted per provider request. +- Each text is limited to 16 KiB. +- Long inputs use overlapping 510-token windows, with 64 content tokens of + overlap. +- ONNX inference batches and total windows are bounded by worker configuration. +- Requests above the provider bounds return an error; the PII component then + fails closed for the affected batch. +- `max_latency_ms` is one total budget for all provider batches selected from + one payload. +- CPU inference is serialized per worker process. Host deadlines cancel the + RPC, while already-running native inference is allowed to finish before its + admission slot is released. +- Use a `max_latency_ms` of at least 5000 when the selected payload can approach + the 64 KiB provider-request limit. Smaller content-only payloads normally + complete much faster. Benchmark representative inputs on deployment + hardware before lowering the deadline. + +The default model supports English, Spanish, French, German, Italian, +Portuguese, and Dutch. Its model card documents weak recall for non-Latin +scripts and government identifiers. Do not treat this detector as a complete +security boundary. + +## Attribution + +See [THIRD_PARTY_NOTICES.md](./THIRD_PARTY_NOTICES.md). The model is downloaded +only by the explicit prefetch command and is not redistributed by this package. diff --git a/crates/pii-redaction/providers/rampart/THIRD_PARTY_NOTICES.md b/crates/pii-redaction/providers/rampart/THIRD_PARTY_NOTICES.md new file mode 100644 index 000000000..248eef265 --- /dev/null +++ b/crates/pii-redaction/providers/rampart/THIRD_PARTY_NOTICES.md @@ -0,0 +1,18 @@ + + +# Third-Party Notices + +This optional provider downloads and executes **Rampart**, published by +National Design Studio at +[`nationaldesignstudio/rampart`](https://huggingface.co/nationaldesignstudio/rampart). +The model and its training-data attribution are published under the +[Creative Commons Attribution 4.0 International +license](https://creativecommons.org/licenses/by/4.0/). + +The default provider configuration selects model revision +`b1993e4e68b082835b80ffc65acc03325ea2e501`. Model files are downloaded to the +operator's Hugging Face cache and are not distributed in the NeMo Relay source +or Python package. diff --git a/crates/pii-redaction/providers/rampart/config.schema.json b/crates/pii-redaction/providers/rampart/config.schema.json new file mode 100644 index 000000000..43fac3687 --- /dev/null +++ b/crates/pii-redaction/providers/rampart/config.schema.json @@ -0,0 +1,54 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "title": "NeMo Relay Rampart PII Provider", + "type": "object", + "additionalProperties": false, + "properties": { + "model_id": { + "type": "string", + "minLength": 1, + "maxLength": 1024, + "default": "nationaldesignstudio/rampart", + "description": "Pinned Rampart repository identifier or an explicit local directory containing the same verified snapshot." + }, + "revision": { + "type": "string", + "const": "b1993e4e68b082835b80ffc65acc03325ea2e501", + "default": "b1993e4e68b082835b80ffc65acc03325ea2e501" + }, + "cache_dir": { + "type": "string", + "minLength": 1, + "maxLength": 4096 + }, + "local_files_only": { + "type": "boolean", + "const": true, + "default": true + }, + "max_windows_per_request": { + "type": "integer", + "minimum": 1, + "maximum": 512, + "default": 128 + }, + "inference_batch_size": { + "type": "integer", + "minimum": 1, + "maximum": 64, + "default": 16, + "description": "Maximum number of token windows per ONNX call. The worker reduces this automatically to bound padded token volume." + }, + "max_pending_requests": { + "type": "integer", + "minimum": 1, + "maximum": 64, + "default": 8 + }, + "intra_op_threads": { + "type": "integer", + "minimum": 1, + "maximum": 64 + } + } +} diff --git a/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/__init__.py b/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/__init__.py new file mode 100644 index 000000000..027d78593 --- /dev/null +++ b/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/__init__.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Rampart inference provider for the NeMo Relay PII component.""" + +from .detector import DEFAULT_MODEL_ID, DEFAULT_MODEL_REVISION, RampartDetector, RampartSettings + +__all__ = [ + "DEFAULT_MODEL_ID", + "DEFAULT_MODEL_REVISION", + "RampartDetector", + "RampartSettings", +] diff --git a/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/detector.py b/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/detector.py new file mode 100644 index 000000000..6dd482516 --- /dev/null +++ b/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/detector.py @@ -0,0 +1,572 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bounded ONNX token-classification adapter for the Rampart PII model.""" + +from __future__ import annotations + +import hashlib +import json +import threading +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Protocol + +import numpy as np # ty: ignore[unresolved-import] +import onnxruntime as ort # ty: ignore[unresolved-import] +from huggingface_hub import snapshot_download # ty: ignore[unresolved-import] +from tokenizers import Tokenizer # ty: ignore[unresolved-import] + +DEFAULT_MODEL_ID = "nationaldesignstudio/rampart" +DEFAULT_MODEL_REVISION = "b1993e4e68b082835b80ffc65acc03325ea2e501" +CONTRACT_VERSION = 1 +MAX_TEXTS_PER_REQUEST = 64 +MAX_TEXT_BYTES = 16 * 1024 +MAX_REQUEST_TEXT_BYTES = 64 * 1024 +MODEL_MAX_TOKENS = 512 +SPECIAL_TOKEN_COUNT = 2 +CONTENT_TOKEN_BUDGET = MODEL_MAX_TOKENS - SPECIAL_TOKEN_COUNT +WINDOW_OVERLAP_TOKENS = 64 +MAX_PADDED_TOKENS_PER_BATCH = MODEL_MAX_TOKENS +MAX_MODEL_REFERENCE_BYTES = 1024 +MAX_DETECTOR_PROFILE_BYTES = 1024 +MAX_CACHE_PATH_BYTES = 4096 + +_MODEL_FILE_SHA256 = { + "config.json": "003b84bbcd489f5e782fe5cad8f3249c3653ec880089abb1ccc398a0d895e3e6", + "onnx/model_q4.onnx": "9f27d24949b0581701071ea5ef522d77ccd3f50c525cc91eac4d265b0fc2afe5", + "special_tokens_map.json": "5d5b662e421ea9fac075174bb0688ee0d9431699900b90662acd44b2a350503a", + "tokenizer.json": "98ade711428b42a1b5343c403a73344535e92de8e19359cdb567ef34da210259", + "tokenizer_config.json": "0088a6f8bcdd4014184fb068b83ebb12896a9db2bb269a71f73de83fef24bceb", +} + + +@dataclass(frozen=True) +class RampartSettings: + """Activation-time settings for one Rampart worker.""" + + model_id: str = DEFAULT_MODEL_ID + revision: str = DEFAULT_MODEL_REVISION + cache_dir: str | None = None + local_files_only: bool = True + max_windows_per_request: int = 128 + inference_batch_size: int = 16 + max_pending_requests: int = 8 + intra_op_threads: int | None = None + + @classmethod + def from_config(cls, config: Any) -> RampartSettings: + """Parse and validate dynamic-plugin configuration.""" + if not isinstance(config, dict): + raise TypeError("plugin config must be a JSON object") + allowed = { + "model_id", + "revision", + "cache_dir", + "local_files_only", + "max_windows_per_request", + "inference_batch_size", + "max_pending_requests", + "intra_op_threads", + } + unknown = sorted(set(config) - allowed) + if unknown: + raise ValueError(f"unknown plugin config field(s): {', '.join(unknown)}") + + model_id = _bounded_string( + config.get("model_id", DEFAULT_MODEL_ID), + "model_id", + MAX_MODEL_REFERENCE_BYTES, + ) + revision = _bounded_string( + config.get("revision", DEFAULT_MODEL_REVISION), + "revision", + MAX_MODEL_REFERENCE_BYTES, + ) + if not _uses_explicit_model_path(model_id) and model_id != DEFAULT_MODEL_ID: + raise ValueError(f"model_id must be {DEFAULT_MODEL_ID!r} or an explicit local directory") + if revision != DEFAULT_MODEL_REVISION: + raise ValueError(f"revision must be the pinned Rampart revision {DEFAULT_MODEL_REVISION!r}") + cache_dir = config.get("cache_dir") + if cache_dir is not None: + cache_dir = _bounded_string(cache_dir, "cache_dir", MAX_CACHE_PATH_BYTES) + local_files_only = config.get("local_files_only", True) + if not isinstance(local_files_only, bool): + raise TypeError("local_files_only must be a boolean") + if not local_files_only: + raise ValueError("local_files_only must remain true; prefetch the pinned model before enabling the plugin") + + return cls( + model_id=model_id, + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + max_windows_per_request=_bounded_integer( + config.get("max_windows_per_request", 128), + "max_windows_per_request", + 1, + 512, + ), + inference_batch_size=_bounded_integer( + config.get("inference_batch_size", 16), + "inference_batch_size", + 1, + 64, + ), + max_pending_requests=_bounded_integer( + config.get("max_pending_requests", 8), + "max_pending_requests", + 1, + 64, + ), + intra_op_threads=_optional_bounded_integer(config.get("intra_op_threads"), "intra_op_threads", 1, 64), + ) + + +@dataclass(frozen=True) +class _InputText: + text_id: int + text: str + + +@dataclass(frozen=True) +class _Window: + text_id: int + input_ids: tuple[int, ...] + token_type_ids: tuple[int, ...] + offsets: tuple[tuple[int, int] | None, ...] + + +@dataclass(frozen=True) +class _Span: + start: int + end: int + label: str + score: float + + +class _Tokenizer(Protocol): + def encode(self, sequence: str, add_special_tokens: bool = True) -> Any: ... + + def token_to_id(self, token: str) -> int | None: ... + + +class _Session(Protocol): + def get_inputs(self) -> list[Any]: ... + + def get_outputs(self) -> list[Any]: ... + + def run(self, output_names: list[str], input_feed: dict[str, np.ndarray[Any, Any]]) -> list[Any]: ... + + +class RampartDetector: + """Load one Rampart model and perform bounded, serialized inference.""" + + def __init__( + self, + settings: RampartSettings, + tokenizer: _Tokenizer, + session: _Session, + labels: dict[int, str], + ) -> None: + self.settings = settings + self._tokenizer = tokenizer + self._session = session + self._labels = labels + if set(labels) != set(range(len(labels))): + raise ValueError("Rampart label IDs must be contiguous from zero") + self._lock = threading.Lock() + self._cls_id = _required_token_id(tokenizer, "[CLS]") + self._sep_id = _required_token_id(tokenizer, "[SEP]") + self._pad_id = _required_token_id(tokenizer, "[PAD]") + self._validate_model_contract() + + @classmethod + def load(cls, settings: RampartSettings) -> RampartDetector: + """Resolve model files and initialize an optimized CPU session.""" + model_root = resolve_verified_model_root(settings) + config = json.loads((model_root / "config.json").read_text(encoding="utf-8")) + raw_labels = config.get("id2label") + if not isinstance(raw_labels, dict): + raise ValueError("Rampart config.json must contain an id2label object") + labels = {int(index): str(label) for index, label in raw_labels.items()} + if not labels or labels.get(0) != "O": + raise ValueError("Rampart label map must define label 0 as O") + + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + options.inter_op_num_threads = 1 + if settings.intra_op_threads is not None: + options.intra_op_num_threads = settings.intra_op_threads + session = ort.InferenceSession( + str(model_root / "onnx" / "model_q4.onnx"), + sess_options=options, + providers=["CPUExecutionProvider"], + ) + tokenizer = Tokenizer.from_file(str(model_root / "tokenizer.json")) + detector = cls(settings, tokenizer, session, labels) + detector._detect_texts([_InputText(0, "warmup")]) + return detector + + def detect_request(self, request: Any) -> dict[str, Any]: + """Validate one provider request and return versioned UTF-8 spans.""" + texts, requested_model, profile = _parse_request(request) + if requested_model is not None and requested_model != DEFAULT_MODEL_ID: + raise ValueError(f"request model_id {requested_model!r} does not match loaded model {DEFAULT_MODEL_ID!r}") + if profile not in (None, "default"): + raise ValueError(f"unsupported detector_profile {profile!r}") + + with self._lock: + detections = self._detect_texts(texts) + return { + "version": CONTRACT_VERSION, + "detections": detections, + } + + def _validate_model_contract(self) -> None: + inputs = {entry.name for entry in self._session.get_inputs()} + expected_inputs = {"input_ids", "attention_mask", "token_type_ids"} + if inputs != expected_inputs: + raise ValueError(f"Rampart ONNX inputs must be {sorted(expected_inputs)}, got {sorted(inputs)}") + outputs = self._session.get_outputs() + if len(outputs) != 1 or outputs[0].name != "logits": + raise ValueError("Rampart ONNX model must expose one logits output") + + def _detect_texts(self, texts: list[_InputText]) -> list[dict[str, Any]]: + windows = self._build_windows(texts) + spans_by_text: dict[int, list[_Span]] = {item.text_id: [] for item in texts} + for batch in _inference_batches(windows, self.settings.inference_batch_size): + logits = self._infer(batch) + for window, window_logits in zip(batch, logits, strict=True): + spans_by_text[window.text_id].extend(self._decode_window(window, window_logits)) + + detections = [] + text_by_id = {item.text_id: item.text for item in texts} + for text_id, spans in spans_by_text.items(): + merged_spans = _merge_overlapping_spans(spans) + if not merged_spans: + continue + byte_offsets = _utf8_offsets(text_by_id[text_id]) + for span in merged_spans: + detections.append( + { + "text_id": text_id, + "start_utf8": byte_offsets[span.start], + "end_utf8": byte_offsets[span.end], + "label": span.label, + "score": span.score, + } + ) + return detections + + def _build_windows(self, texts: list[_InputText]) -> list[_Window]: + windows = [] + step = CONTENT_TOKEN_BUDGET - WINDOW_OVERLAP_TOKENS + for item in texts: + inference_text = item.text.replace("-", " ") + encoding = self._tokenizer.encode(inference_text, add_special_tokens=False) + ids = list(encoding.ids) + type_ids = list(encoding.type_ids) + offsets = list(encoding.offsets) + if not (len(ids) == len(type_ids) == len(offsets)): + raise ValueError("tokenizer returned inconsistent token metadata") + if any(start < 0 or end < start or end > len(inference_text) for start, end in offsets): + raise ValueError("tokenizer returned invalid character offsets") + for start in range(0, len(ids), step): + end = min(start + CONTENT_TOKEN_BUDGET, len(ids)) + windows.append( + _Window( + text_id=item.text_id, + input_ids=(self._cls_id, *ids[start:end], self._sep_id), + token_type_ids=(0, *type_ids[start:end], 0), + offsets=(None, *offsets[start:end], None), + ) + ) + if len(windows) > self.settings.max_windows_per_request: + raise ValueError( + f"request exceeded max_windows_per_request={self.settings.max_windows_per_request}" + ) + if end == len(ids): + break + return windows + + def _infer(self, windows: list[_Window]) -> np.ndarray[Any, np.dtype[np.float32]]: + max_length = max(len(window.input_ids) for window in windows) + shape = (len(windows), max_length) + input_ids = np.full(shape, self._pad_id, dtype=np.int64) + attention_mask = np.zeros(shape, dtype=np.int64) + token_type_ids = np.zeros(shape, dtype=np.int64) + for index, window in enumerate(windows): + length = len(window.input_ids) + input_ids[index, :length] = window.input_ids + attention_mask[index, :length] = 1 + token_type_ids[index, :length] = window.token_type_ids + result = self._session.run( + ["logits"], + { + "input_ids": input_ids, + "attention_mask": attention_mask, + "token_type_ids": token_type_ids, + }, + ) + logits = np.asarray(result[0], dtype=np.float32) + expected = (len(windows), max_length, len(self._labels)) + if logits.shape != expected: + raise ValueError(f"Rampart logits shape must be {expected}, got {logits.shape}") + if not np.isfinite(logits).all(): + raise ValueError("Rampart logits must contain only finite values") + return logits + + def _decode_window(self, window: _Window, logits: np.ndarray[Any, Any]) -> list[_Span]: + label_ids = np.argmax(logits, axis=-1) + maxima = np.max(logits, axis=-1) + scores = 1.0 / np.exp(logits - maxima[:, None]).sum(axis=-1) + spans = [] + current_label: str | None = None + current_start = 0 + current_end = 0 + current_score = 0.0 + current_count = 0 + + def finish() -> None: + nonlocal current_label, current_start, current_end, current_score, current_count + if current_label is not None: + score = current_score / current_count + spans.append(_Span(current_start, current_end, current_label, score)) + current_label = None + current_start = 0 + current_end = 0 + current_score = 0.0 + current_count = 0 + + for index, offset in enumerate(window.offsets): + if offset is None or offset[0] >= offset[1]: + finish() + continue + raw_label = self._labels.get(int(label_ids[index])) + score = float(scores[index]) + prefix, label = _split_bio_label(raw_label) + if label is None: + finish() + continue + if current_label is None or prefix == "B" or label != current_label: + finish() + current_label = label + current_start = offset[0] + current_end = offset[1] + current_score = score + current_count = 1 + else: + current_end = max(current_end, offset[1]) + current_score += score + current_count += 1 + finish() + return spans + + +def _inference_batches(windows: list[_Window], max_batch_size: int) -> Iterator[list[_Window]]: + ordered = sorted(windows, key=lambda window: len(window.input_ids)) + batch: list[_Window] = [] + max_tokens = 0 + for window in ordered: + next_max_tokens = max(max_tokens, len(window.input_ids)) + padded_tokens = (len(batch) + 1) * next_max_tokens + if batch and (len(batch) >= max_batch_size or padded_tokens > MAX_PADDED_TOKENS_PER_BATCH): + yield batch + batch = [] + max_tokens = 0 + batch.append(window) + max_tokens = max(max_tokens, len(window.input_ids)) + if batch: + yield batch + + +def _resolve_model_root(settings: RampartSettings) -> Path: + local_path = _explicit_model_root(settings.model_id) + if local_path is not None: + return local_path + resolved = snapshot_download( + settings.model_id, + revision=settings.revision, + cache_dir=settings.cache_dir, + allow_patterns=list(_MODEL_FILE_SHA256), + local_files_only=True, + ) + return Path(resolved) + + +def resolve_verified_model_root(settings: RampartSettings) -> Path: + """Resolve and verify the pinned model assets without loading ONNX Runtime.""" + model_root = _resolve_model_root(settings) + _verify_model_files(model_root) + return model_root + + +def prefetch_verified_model(cache_dir: str | None = None) -> Path: + """Download and verify the pinned Rampart assets outside plugin activation.""" + if cache_dir is not None: + cache_dir = _bounded_string(cache_dir, "cache_dir", MAX_CACHE_PATH_BYTES) + model_root = Path( + snapshot_download( + DEFAULT_MODEL_ID, + revision=DEFAULT_MODEL_REVISION, + cache_dir=cache_dir, + allow_patterns=list(_MODEL_FILE_SHA256), + local_files_only=False, + ) + ) + _verify_model_files(model_root) + return model_root + + +def _verify_model_files( + model_root: Path, + expected: Mapping[str, str] = _MODEL_FILE_SHA256, +) -> None: + for relative_path, expected_sha256 in expected.items(): + path = model_root / relative_path + if not path.is_file(): + raise ValueError(f"Rampart model is missing required file {relative_path!r}") + with path.open("rb") as model_file: + digest = hashlib.file_digest(model_file, "sha256").hexdigest() + if digest != expected_sha256: + raise ValueError(f"Rampart model file {relative_path!r} failed SHA-256 verification") + + +def _explicit_model_root(model_id: str) -> Path | None: + candidate = Path(model_id).expanduser() + if not _uses_explicit_model_path(model_id): + return None + if not candidate.is_dir(): + raise ValueError(f"local Rampart model directory does not exist: {model_id}") + return candidate.resolve() + + +def _uses_explicit_model_path(model_id: str) -> bool: + return Path(model_id).expanduser().is_absolute() or model_id.startswith(("./", "../", ".\\", "..\\", "~/", "~\\")) + + +def _parse_request(request: Any) -> tuple[list[_InputText], str | None, str | None]: + if not isinstance(request, dict): + raise TypeError("local-model request must be a JSON object") + allowed = {"version", "model_id", "detector_profile", "texts"} + unknown = sorted(set(request) - allowed) + if unknown: + raise ValueError(f"unknown local-model request field(s): {', '.join(unknown)}") + version = request.get("version") + if isinstance(version, bool) or version != CONTRACT_VERSION: + raise ValueError(f"local-model request version must be {CONTRACT_VERSION}") + model_id = request.get("model_id") + if model_id is not None: + model_id = _bounded_string(model_id, "model_id", MAX_MODEL_REFERENCE_BYTES) + profile = request.get("detector_profile") + if profile is not None: + profile = _bounded_string(profile, "detector_profile", MAX_DETECTOR_PROFILE_BYTES) + raw_texts = request.get("texts") + if not isinstance(raw_texts, list) or not raw_texts: + raise TypeError("texts must be a non-empty array") + if len(raw_texts) > MAX_TEXTS_PER_REQUEST: + raise ValueError(f"texts must contain at most {MAX_TEXTS_PER_REQUEST} items") + + texts = [] + seen_ids = set() + total_bytes = 0 + for item in raw_texts: + if not isinstance(item, dict) or set(item) != {"id", "text"}: + raise TypeError("each texts item must contain exactly id and text") + text_id = item["id"] + text = item["text"] + if isinstance(text_id, bool) or not isinstance(text_id, int) or not 0 <= text_id <= 2**32 - 1: + raise TypeError("text id must be an unsigned 32-bit integer") + if text_id in seen_ids: + raise ValueError(f"duplicate text id {text_id}") + if not isinstance(text, str): + raise TypeError("text must be a string") + text_bytes = len(text.encode("utf-8")) + if text_bytes > MAX_TEXT_BYTES: + raise ValueError(f"text {text_id} exceeds {MAX_TEXT_BYTES} UTF-8 bytes") + total_bytes += text_bytes + if total_bytes > MAX_REQUEST_TEXT_BYTES: + raise ValueError(f"request text exceeds {MAX_REQUEST_TEXT_BYTES} UTF-8 bytes") + seen_ids.add(text_id) + texts.append(_InputText(text_id, text)) + return texts, model_id, profile + + +def _split_bio_label(raw_label: str | None) -> tuple[str | None, str | None]: + if raw_label is None or raw_label == "O": + return None, None + if raw_label.startswith(("B-", "I-")) and len(raw_label) > 2: + return raw_label[0], raw_label[2:].upper() + return "B", raw_label.upper() + + +def _merge_overlapping_spans(spans: list[_Span]) -> list[_Span]: + merged: list[_Span] = [] + for span in sorted(spans, key=lambda item: (item.start, -item.end, -item.score, item.label)): + if ( + not merged + or span.start > merged[-1].end + or (span.start == merged[-1].end and span.label != merged[-1].label) + ): + merged.append(span) + continue + previous = merged[-1] + winner = _preferred_span(previous, span) + merged[-1] = _Span( + start=min(previous.start, span.start), + end=max(previous.end, span.end), + label=winner.label, + score=max(previous.score, span.score), + ) + return merged + + +def _preferred_span(left: _Span, right: _Span) -> _Span: + left_key = (left.score, left.end - left.start, left.label) + right_key = (right.score, right.end - right.start, right.label) + return left if left_key >= right_key else right + + +def _utf8_offsets(text: str) -> list[int]: + offsets = [0] + total = 0 + for character in text: + total += len(character.encode("utf-8")) + offsets.append(total) + return offsets + + +def _required_token_id(tokenizer: _Tokenizer, token: str) -> int: + token_id = tokenizer.token_to_id(token) + if token_id is None: + raise ValueError(f"tokenizer is missing required token {token}") + return token_id + + +def _nonempty_string(value: Any, name: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise TypeError(f"{name} must be a non-empty string") + return value + + +def _bounded_string(value: Any, name: str, maximum_bytes: int) -> str: + value = _nonempty_string(value, name) + if len(value.encode("utf-8")) > maximum_bytes: + raise ValueError(f"{name} must not exceed {maximum_bytes} UTF-8 bytes") + return value + + +def _bounded_integer(value: Any, name: str, minimum: int, maximum: int) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer") + if not minimum <= value <= maximum: + raise ValueError(f"{name} must be between {minimum} and {maximum}") + return value + + +def _optional_bounded_integer(value: Any, name: str, minimum: int, maximum: int) -> int | None: + if value is None: + return None + return _bounded_integer(value, name, minimum, maximum) diff --git a/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/prefetch.py b/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/prefetch.py new file mode 100644 index 000000000..f75d6b213 --- /dev/null +++ b/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/prefetch.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Prefetch the pinned Rampart model before enabling the worker.""" + +from __future__ import annotations + +import argparse +from collections.abc import Sequence + +from .detector import prefetch_verified_model + + +def main(argv: Sequence[str] | None = None) -> None: + """Download and verify the model in the configured Hugging Face cache.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--cache-dir", + help="Optional Hugging Face cache directory shared with the worker.", + ) + args = parser.parse_args(argv) + model_root = prefetch_verified_model(args.cache_dir) + print(f"Verified Rampart model at {model_root}") + + +if __name__ == "__main__": + main() diff --git a/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/py.typed b/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/worker.py b/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/worker.py new file mode 100644 index 000000000..d04533e99 --- /dev/null +++ b/crates/pii-redaction/providers/rampart/nemo_relay_pii_rampart/worker.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Manifest entrypoint for the Rampart PII inference provider.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +from nemo_relay_plugin import ConfigDiagnostic, DiagnosticLevel, Json, PluginContext, WorkerPlugin, serve_plugin + +from .detector import RampartDetector, RampartSettings, resolve_verified_model_root + +PII_DETECTION_PROVIDER_CONTRACT = "nemo.relay.pii_detection.v1" + + +class _Admission: + def __init__(self, limit: int) -> None: + self._limit = limit + self._active = 0 + + def acquire(self) -> None: + if self._active >= self._limit: + raise RuntimeError("Rampart provider is at its pending-request limit") + self._active += 1 + + def release(self) -> None: + self._active -= 1 + + +class RampartWorker(WorkerPlugin): + """Expose Rampart inference through the PII component's provider contract.""" + + plugin_id = "nemo_relay.pii_rampart" + + def validate(self, config: Json) -> list[ConfigDiagnostic | dict[str, Any]]: + try: + settings = RampartSettings.from_config(config) + except (TypeError, ValueError) as error: + return [ + ConfigDiagnostic( + level=DiagnosticLevel.ERROR, + code="nemo_relay.pii_rampart.invalid_config", + component=self.plugin_id, + message=str(error), + ) + ] + try: + resolve_verified_model_root(settings) + except Exception: + return [ + ConfigDiagnostic( + level=DiagnosticLevel.ERROR, + code="nemo_relay.pii_rampart.model_unavailable", + component=self.plugin_id, + message=( + "the pinned Rampart model is unavailable or failed integrity " + "verification; prefetch it before enabling the plugin" + ), + ) + ] + return [] + + def register(self, ctx: PluginContext, config: Json) -> None: + settings = RampartSettings.from_config(config) + detector = RampartDetector.load(settings) + admission = _Admission(settings.max_pending_requests) + inference_slot = asyncio.Lock() + + async def detect(request: Json) -> Json: + admission.acquire() + native_started = False + + async def run_native() -> Json: + nonlocal native_started + async with inference_slot: + native_started = True + return await asyncio.to_thread(detector.detect_request, request) + + work = asyncio.create_task(run_native()) + release_on_completion = False + try: + return await asyncio.shield(work) + except asyncio.CancelledError: + release_on_completion = True + if not native_started: + work.cancel() + + def release_after_work(_task: asyncio.Task[Json]) -> None: + try: + _task.exception() + except asyncio.CancelledError: + pass + admission.release() + + work.add_done_callback(release_after_work) + raise + finally: + if not release_on_completion: + admission.release() + + ctx.register_inference_provider( + "detector", + PII_DETECTION_PROVIDER_CONTRACT, + detect, + ) + + +async def main() -> None: + """Start the Relay-managed worker.""" + await serve_plugin(RampartWorker()) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/crates/pii-redaction/providers/rampart/pyproject.toml b/crates/pii-redaction/providers/rampart/pyproject.toml new file mode 100644 index 000000000..c658534cb --- /dev/null +++ b/crates/pii-redaction/providers/rampart/pyproject.toml @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["setuptools>=77"] +build-backend = "setuptools.build_meta" + +[project] +name = "nemo-relay-pii-rampart" +version = "0.1.0" +description = "Optional Rampart inference provider for NeMo Relay PII redaction" +readme = "README.md" +requires-python = ">=3.11" +license = "Apache-2.0" +license-files = ["THIRD_PARTY_NOTICES.md"] +authors = [ + { name = "NVIDIA Corporation & Affiliates" }, +] +dependencies = [ + "huggingface-hub>=0.34,<2", + "nemo-relay-plugin>=0.7,<1.0", + "numpy>=1.26,<3", + "onnxruntime>=1.20,<2", + "tokenizers>=0.21,<1", +] + +[project.scripts] +nemo-relay-pii-rampart-prefetch = "nemo_relay_pii_rampart.prefetch:main" + +[project.optional-dependencies] +test = [ + "pytest>=8", +] + +[tool.setuptools.packages.find] +where = ["."] +include = ["nemo_relay_pii_rampart"] + +[tool.setuptools.package-data] +nemo_relay_pii_rampart = ["py.typed"] + +[tool.ruff] +line-length = 120 +target-version = "py311" + +[tool.ruff.format] +quote-style = "double" + +[tool.ruff.lint] +select = ["E", "F", "W", "I"] + +[tool.ruff.lint.isort] +known-first-party = ["nemo_relay_pii_rampart", "nemo_relay_plugin"] + +[tool.ty.analysis] +# The model dependencies live only in the worker's managed environment. +allowed-unresolved-imports = ["huggingface_hub", "numpy", "onnxruntime", "tokenizers"] diff --git a/crates/pii-redaction/providers/rampart/relay-plugin.toml b/crates/pii-redaction/providers/rampart/relay-plugin.toml new file mode 100644 index 000000000..a83fdaf18 --- /dev/null +++ b/crates/pii-redaction/providers/rampart/relay-plugin.toml @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +manifest_version = 1 + +[plugin] +id = "nemo_relay.pii_rampart" +kind = "worker" + +[compat] +relay = ">=0.7,<1.0" +worker_protocol = "grpc-v1" + +[defaults] +enabled = false + +[capabilities] +items = ["plugin_worker", "config_schema"] + +[config_schema] +path = "config.schema.json" + +[source] +manifest_root = "." +artifact = "nemo_relay_pii_rampart/worker.py" + +[integrity] +sha256 = "sha256:e0ab2677112b8687d0bee32b257373a4057878f437ec8b66a233cb7b02533338" + +[load] +runtime = "python" +entrypoint = "nemo_relay_pii_rampart.worker:main" diff --git a/crates/pii-redaction/providers/rampart/tests/test_detector.py b/crates/pii-redaction/providers/rampart/tests/test_detector.py new file mode 100644 index 000000000..d1d75a575 --- /dev/null +++ b/crates/pii-redaction/providers/rampart/tests/test_detector.py @@ -0,0 +1,339 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np # ty: ignore[unresolved-import] +import pytest + +import nemo_relay_pii_rampart.detector as detector_module +from nemo_relay_pii_rampart.detector import ( + DEFAULT_MODEL_ID, + RampartDetector, + RampartSettings, + _explicit_model_root, + _merge_overlapping_spans, + _parse_request, + _Span, + _verify_model_files, + prefetch_verified_model, +) + + +class FakeTokenizer: + _tokens = {"[PAD]": 0, "[CLS]": 2, "[SEP]": 3} + + def token_to_id(self, token: str) -> int | None: + return self._tokens.get(token) + + def encode(self, sequence: str, add_special_tokens: bool = True) -> SimpleNamespace: + del add_special_tokens + words = sequence.split() + ids = [] + offsets = [] + cursor = 0 + for index, word in enumerate(words): + start = sequence.index(word, cursor) + end = start + len(word) + ids.append(10 + index) + offsets.append((start, end)) + cursor = end + return SimpleNamespace(ids=ids, type_ids=[0] * len(ids), offsets=offsets) + + +class FakeSession: + def __init__(self, label_ids: list[int], scores: list[float]) -> None: + self._label_ids = label_ids + self._scores = scores + + def get_inputs(self) -> list[SimpleNamespace]: + return [SimpleNamespace(name=name) for name in ("input_ids", "attention_mask", "token_type_ids")] + + def get_outputs(self) -> list[SimpleNamespace]: + return [SimpleNamespace(name="logits")] + + def run(self, output_names: list[str], input_feed: dict[str, np.ndarray]) -> list[np.ndarray]: + assert output_names == ["logits"] + shape = (*input_feed["input_ids"].shape, 5) + logits = np.full(shape, -10.0, dtype=np.float32) + logits[:, :, 0] = 10.0 + for token_index, (label_id, score) in enumerate(zip(self._label_ids, self._scores, strict=True), start=1): + logits[:, token_index, 0] = 0.0 + logits[:, token_index, label_id] = np.log(score / (1.0 - score) * 4.0) + return [logits] + + +class NonFiniteSession(FakeSession): + def run(self, output_names: list[str], input_feed: dict[str, np.ndarray]) -> list[np.ndarray]: + logits = super().run(output_names, input_feed)[0] + logits[0, 0, 0] = np.nan + return [logits] + + +class RecordingSession(FakeSession): + def __init__(self) -> None: + super().__init__([], []) + self.shapes: list[tuple[int, int]] = [] + + def run(self, output_names: list[str], input_feed: dict[str, np.ndarray]) -> list[np.ndarray]: + self.shapes.append(input_feed["input_ids"].shape) + return super().run(output_names, input_feed) + + +def detector( + label_ids: list[int], + scores: list[float], + *, + max_windows_per_request: int | None = None, +) -> RampartDetector: + config = {} if max_windows_per_request is None else {"max_windows_per_request": max_windows_per_request} + return RampartDetector( + RampartSettings.from_config(config), + FakeTokenizer(), + FakeSession(label_ids, scores), + {0: "O", 1: "B-GIVEN_NAME", 2: "I-GIVEN_NAME", 3: "B-CITY", 4: "I-CITY"}, + ) + + +def test_settings_validate_unknown_and_bounded_fields() -> None: + settings = RampartSettings.from_config({}) + assert settings.model_id == DEFAULT_MODEL_ID + assert settings.local_files_only is True + with pytest.raises(ValueError, match="unknown plugin config"): + RampartSettings.from_config({"surprise": True}) + with pytest.raises(TypeError, match="max_pending_requests"): + RampartSettings.from_config({"max_pending_requests": True}) + with pytest.raises(ValueError, match="model_id"): + RampartSettings.from_config({"model_id": "x" * 1025}) + with pytest.raises(ValueError, match="explicit local directory"): + RampartSettings.from_config({"model_id": "other/model"}) + with pytest.raises(ValueError, match="pinned Rampart revision"): + RampartSettings.from_config({"revision": "main"}) + with pytest.raises(ValueError, match="prefetch"): + RampartSettings.from_config({"local_files_only": False}) + + +def test_model_files_are_verified_before_loading(tmp_path: Path) -> None: + model_file = tmp_path / "model.bin" + model_file.write_bytes(b"trusted model") + expected = {"model.bin": hashlib.sha256(b"trusted model").hexdigest()} + + _verify_model_files(tmp_path, expected) + model_file.write_bytes(b"modified model") + with pytest.raises(ValueError, match="SHA-256 verification"): + _verify_model_files(tmp_path, expected) + + +def test_model_file_verification_rejects_missing_files(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="missing required file"): + _verify_model_files(tmp_path, {"missing.bin": "0" * 64}) + + +def test_prefetch_uses_the_pinned_snapshot_and_verifies_it( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + observed: dict[str, Any] = {} + + def snapshot(model_id: str, **kwargs: Any) -> str: + observed["model_id"] = model_id + observed.update(kwargs) + return str(tmp_path) + + monkeypatch.setattr(detector_module, "snapshot_download", snapshot) + monkeypatch.setattr( + detector_module, "_verify_model_files", lambda model_root: observed.setdefault("root", model_root) + ) + + assert prefetch_verified_model("/tmp/rampart-cache") == tmp_path + assert observed["model_id"] == DEFAULT_MODEL_ID + assert observed["revision"] == detector_module.DEFAULT_MODEL_REVISION + assert observed["local_files_only"] is False + assert observed["cache_dir"] == "/tmp/rampart-cache" + assert observed["root"] == tmp_path + + +def test_local_model_directories_require_explicit_path_syntax( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_root = tmp_path / "model" + model_root.mkdir() + assert _explicit_model_root(str(model_root)) == model_root.resolve() + settings = RampartSettings.from_config({"model_id": str(model_root)}) + assert settings.model_id == str(model_root) + + shadow = tmp_path / "nationaldesignstudio" / "rampart" + shadow.mkdir(parents=True) + monkeypatch.chdir(tmp_path) + assert _explicit_model_root(DEFAULT_MODEL_ID) is None + with pytest.raises(ValueError, match="does not exist"): + _explicit_model_root("./missing") + + +def test_request_validation_rejects_duplicate_ids_and_byte_overflow() -> None: + with pytest.raises(ValueError, match="duplicate text id"): + _parse_request( + { + "version": 1, + "texts": [ + {"id": 0, "text": "one"}, + {"id": 0, "text": "two"}, + ], + } + ) + with pytest.raises(ValueError, match="UTF-8 bytes"): + _parse_request({"version": 1, "texts": [{"id": 0, "text": "é" * 9000}]}) + with pytest.raises(ValueError, match="model_id"): + _parse_request( + { + "version": 1, + "model_id": "x" * 1025, + "texts": [{"id": 0, "text": ""}], + } + ) + with pytest.raises(ValueError, match="detector_profile"): + _parse_request( + { + "version": 1, + "detector_profile": "x" * 1025, + "texts": [{"id": 0, "text": ""}], + } + ) + + +def test_no_detection_does_not_allocate_utf8_offset_tables(monkeypatch: pytest.MonkeyPatch) -> None: + def unexpected_offsets(_text: str) -> list[int]: + raise AssertionError("UTF-8 offsets should be lazy when no spans were detected") + + monkeypatch.setattr(detector_module, "_utf8_offsets", unexpected_offsets) + assert detector([], []).detect_request( + { + "version": 1, + "texts": [{"id": 0, "text": "no private values"}], + } + ) == {"version": 1, "detections": []} + + +def test_detector_returns_utf8_byte_offsets_and_model_labels() -> None: + value = detector([1, 2], [0.99, 0.98]).detect_request({"version": 1, "texts": [{"id": 7, "text": "José Rivera"}]}) + assert value["version"] == 1 + assert len(value["detections"]) == 1 + detection = value["detections"][0] + assert detection["text_id"] == 7 + assert detection["start_utf8"] == 0 + assert detection["end_utf8"] == len("José Rivera".encode()) + assert detection["label"] == "GIVEN_NAME" + assert 0.9 <= detection["score"] <= 1.0 + + city = detector([3, 4], [0.99, 0.98]).detect_request({"version": 1, "texts": [{"id": 0, "text": "New York"}]}) + assert city["detections"][0]["label"] == "CITY" + + +def test_detector_rejects_model_and_profile_mismatch() -> None: + current = detector([], []) + with pytest.raises(ValueError, match="does not match loaded model"): + current.detect_request( + { + "version": 1, + "model_id": "other/model", + "texts": [{"id": 0, "text": ""}], + } + ) + with pytest.raises(ValueError, match="unsupported detector_profile"): + current.detect_request( + { + "version": 1, + "detector_profile": "strict", + "texts": [{"id": 0, "text": ""}], + } + ) + + +def test_local_model_path_keeps_the_logical_model_identity(tmp_path: Path) -> None: + current = RampartDetector( + RampartSettings.from_config({"model_id": str(tmp_path)}), + FakeTokenizer(), + FakeSession([], []), + {0: "O", 1: "B-GIVEN_NAME", 2: "I-GIVEN_NAME", 3: "B-CITY", 4: "I-CITY"}, + ) + result = current.detect_request( + { + "version": 1, + "model_id": DEFAULT_MODEL_ID, + "texts": [{"id": 0, "text": ""}], + } + ) + assert result == {"version": 1, "detections": []} + + +def test_overlapping_window_spans_are_coalesced() -> None: + spans = _merge_overlapping_spans( + [ + _Span(0, 10, "GIVEN_NAME", 0.8), + _Span(5, 12, "SURNAME", 0.9), + _Span(20, 24, "PHONE", 0.7), + _Span(24, 28, "PHONE", 0.8), + _Span(28, 30, "TAX_ID", 0.9), + ] + ) + assert spans == [ + _Span(0, 12, "SURNAME", 0.9), + _Span(20, 28, "PHONE", 0.8), + _Span(28, 30, "TAX_ID", 0.9), + ] + + +def test_request_window_limit_is_enforced_before_inference() -> None: + current = detector([], [], max_windows_per_request=1) + text = " ".join(f"word{index}" for index in range(600)) + with pytest.raises(ValueError, match="max_windows_per_request"): + current.detect_request({"version": 1, "texts": [{"id": 0, "text": text}]}) + + +def test_inference_batches_short_windows_and_isolates_full_windows() -> None: + session = RecordingSession() + current = RampartDetector( + RampartSettings.from_config({"inference_batch_size": 16}), + FakeTokenizer(), + session, + {0: "O", 1: "B-GIVEN_NAME", 2: "I-GIVEN_NAME", 3: "B-CITY", 4: "I-CITY"}, + ) + full_window = " ".join("word" for _ in range(detector_module.CONTENT_TOKEN_BUDGET)) + texts = [{"id": index, "text": "short"} for index in range(16)] + texts.extend( + [ + {"id": 16, "text": full_window}, + {"id": 17, "text": full_window}, + ] + ) + + assert current.detect_request({"version": 1, "texts": texts}) == { + "version": 1, + "detections": [], + } + assert session.shapes == [(16, 3), (1, detector_module.MODEL_MAX_TOKENS), (1, detector_module.MODEL_MAX_TOKENS)] + + +def test_detector_rejects_invalid_model_outputs() -> None: + with pytest.raises(ValueError, match="contiguous"): + RampartDetector( + RampartSettings(), + FakeTokenizer(), + FakeSession([], []), + {0: "O", 2: "B-GIVEN_NAME"}, + ) + + current = RampartDetector( + RampartSettings(), + FakeTokenizer(), + NonFiniteSession([], []), + {0: "O", 1: "B-GIVEN_NAME", 2: "I-GIVEN_NAME", 3: "B-CITY", 4: "I-CITY"}, + ) + with pytest.raises(ValueError, match="finite"): + current.detect_request({"version": 1, "texts": [{"id": 0, "text": "hello"}]}) diff --git a/crates/pii-redaction/providers/rampart/tests/test_worker.py b/crates/pii-redaction/providers/rampart/tests/test_worker.py new file mode 100644 index 000000000..e7e7fc0b7 --- /dev/null +++ b/crates/pii-redaction/providers/rampart/tests/test_worker.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +import threading +from typing import Any, cast + +import pytest + +import nemo_relay_pii_rampart.worker as worker_module +from nemo_relay_pii_rampart.worker import RampartWorker +from nemo_relay_plugin import ConfigDiagnostic, PluginContext + + +class FakeContext: + def __init__(self) -> None: + self.callback: Any = None + + def register_inference_provider(self, name: str, contract: str, callback: Any) -> None: + assert name == "detector" + assert contract == "nemo.relay.pii_detection.v1" + self.callback = callback + + +class FakeDetector: + def __init__(self, started: threading.Event | None = None, release: threading.Event | None = None) -> None: + self.started = started + self.release = release + + def detect_request(self, request: Any) -> dict[str, Any]: + if self.started is not None: + self.started.set() + if self.release is not None: + self.release.wait(timeout=5) + return {"version": 1, "detections": [], "echo": request} + + +class FailingDetector: + def detect_request(self, request: Any) -> dict[str, Any]: + del request + raise RuntimeError("detector failed") + + +def test_worker_validation_reports_invalid_config(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(worker_module, "resolve_verified_model_root", lambda _settings: None) + worker = RampartWorker() + assert worker.validate({}) == [] + diagnostics = worker.validate({"max_pending_requests": "many"}) + assert len(diagnostics) == 1 + diagnostic = diagnostics[0] + assert isinstance(diagnostic, ConfigDiagnostic) + assert diagnostic.code == "nemo_relay.pii_rampart.invalid_config" + diagnostics = worker.validate({"local_files_only": False}) + assert len(diagnostics) == 1 + diagnostic = diagnostics[0] + assert isinstance(diagnostic, ConfigDiagnostic) + assert diagnostic.code == "nemo_relay.pii_rampart.invalid_config" + + +def test_worker_validation_reports_bounded_model_readiness_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fail(_settings: Any) -> None: + raise ValueError("/sensitive/cache/path/model_q4.onnx is corrupt") + + monkeypatch.setattr(worker_module, "resolve_verified_model_root", fail) + + diagnostics = RampartWorker().validate({}) + + assert len(diagnostics) == 1 + diagnostic = diagnostics[0] + assert isinstance(diagnostic, ConfigDiagnostic) + assert diagnostic.code == "nemo_relay.pii_rampart.model_unavailable" + assert "prefetch" in diagnostic.message + assert "/sensitive" not in diagnostic.message + + +def test_worker_registers_async_provider(monkeypatch: pytest.MonkeyPatch) -> None: + fake = FakeDetector() + monkeypatch.setattr(worker_module.RampartDetector, "load", lambda _settings: fake) + context = FakeContext() + RampartWorker().register(cast(PluginContext, context), {}) + + request = {"version": 1, "texts": [{"id": 0, "text": "hello"}]} + assert asyncio.run(context.callback(request)) == { + "version": 1, + "detections": [], + "echo": request, + } + + +def test_cancelled_callback_holds_admission_until_native_work_finishes(monkeypatch: pytest.MonkeyPatch) -> None: + started = threading.Event() + release = threading.Event() + fake = FakeDetector(started, release) + monkeypatch.setattr(worker_module.RampartDetector, "load", lambda _settings: fake) + context = FakeContext() + RampartWorker().register(cast(PluginContext, context), {"max_pending_requests": 1}) + + async def exercise() -> None: + first = asyncio.create_task(context.callback({"version": 1, "texts": [{"id": 0, "text": "one"}]})) + assert await asyncio.to_thread(started.wait, 1) + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + with pytest.raises(RuntimeError, match="pending-request limit"): + await context.callback({"version": 1, "texts": [{"id": 1, "text": "two"}]}) + + release.set() + for _ in range(100): + await asyncio.sleep(0.01) + try: + result = await context.callback({"version": 1, "texts": [{"id": 2, "text": "three"}]}) + except RuntimeError: + continue + assert result["version"] == 1 + return + pytest.fail("provider admission was not released after native work completed") + + asyncio.run(exercise()) + + +def test_cancelled_queued_callback_does_not_run_native_inference(monkeypatch: pytest.MonkeyPatch) -> None: + started = threading.Event() + release = threading.Event() + requests: list[int] = [] + + class RecordingDetector: + def detect_request(self, request: Any) -> dict[str, Any]: + requests.append(request["texts"][0]["id"]) + started.set() + release.wait(timeout=5) + return {"version": 1, "detections": []} + + monkeypatch.setattr(worker_module.RampartDetector, "load", lambda _settings: RecordingDetector()) + context = FakeContext() + RampartWorker().register(cast(PluginContext, context), {"max_pending_requests": 2}) + + async def exercise() -> None: + first = asyncio.create_task(context.callback({"version": 1, "texts": [{"id": 0, "text": "one"}]})) + assert await asyncio.to_thread(started.wait, 1) + second = asyncio.create_task(context.callback({"version": 1, "texts": [{"id": 1, "text": "two"}]})) + await asyncio.sleep(0) + second.cancel() + with pytest.raises(asyncio.CancelledError): + await second + + release.set() + await first + await context.callback({"version": 1, "texts": [{"id": 2, "text": "three"}]}) + + asyncio.run(exercise()) + assert requests == [0, 2] + + +def test_detector_failure_releases_admission(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(worker_module.RampartDetector, "load", lambda _settings: FailingDetector()) + context = FakeContext() + RampartWorker().register(cast(PluginContext, context), {"max_pending_requests": 1}) + + async def exercise() -> None: + for text_id in range(2): + with pytest.raises(RuntimeError, match="detector failed"): + await context.callback( + { + "version": 1, + "texts": [{"id": text_id, "text": "private"}], + } + ) + + asyncio.run(exercise()) diff --git a/crates/pii-redaction/src/builtin.rs b/crates/pii-redaction/src/builtin.rs index 6e6d1dc5c..292da0430 100644 --- a/crates/pii-redaction/src/builtin.rs +++ b/crates/pii-redaction/src/builtin.rs @@ -534,6 +534,10 @@ pub(super) fn llm_sanitize_request_callback( request.content = backend.sanitize_json_preorder_dfs(request.content); return Some(request); } + if matches!(context.codec(), LlmCodecIdentity::None) && backend.legacy_surface.is_none() { + request.content = backend.sanitize_json_preorder_dfs(request.content); + return Some(request); + } let resolved = context.resolve_codec(); let fallback = if resolved.is_none() { backend @@ -568,6 +572,9 @@ pub(super) fn llm_sanitize_response_callback( if backend.target_paths.is_empty() { return Some(backend.sanitize_json_preorder_dfs(payload)); } + if matches!(context.codec(), LlmCodecIdentity::None) && backend.legacy_surface.is_none() { + return Some(backend.sanitize_json_preorder_dfs(payload)); + } if matches!(context.codec(), LlmCodecIdentity::None) && !backend.uses_compatible_legacy_response_codec(&payload) { @@ -696,7 +703,7 @@ fn remove_sanitized_json_pointer_value(value: &mut Json, segments: &[String]) -> } } -fn render_json_pointer_path(path_segments: &[String]) -> String { +pub(super) fn render_json_pointer_path(path_segments: &[String]) -> String { if path_segments.is_empty() { return String::new(); } @@ -708,7 +715,7 @@ fn render_json_pointer_path(path_segments: &[String]) -> String { rendered } -fn escape_json_pointer_segment(segment: &str) -> String { +pub(super) fn escape_json_pointer_segment(segment: &str) -> String { segment.replace('~', "~0").replace('/', "~1") } diff --git a/crates/pii-redaction/src/component.rs b/crates/pii-redaction/src/component.rs index 2f6b9bfd3..bddd024c2 100644 --- a/crates/pii-redaction/src/component.rs +++ b/crates/pii-redaction/src/component.rs @@ -24,11 +24,45 @@ use super::builtin::{ #[cfg(test)] pub(crate) use super::builtin::{hex_sha256, mask_text}; use super::detectors::{detector_regex_pattern, supported_detector_summary}; -use super::local::register_local_backend; -pub use super::local::{clear_local_backend_provider, register_local_backend_provider}; +use super::local::{register_local_backend, validate_local_backend_config}; /// The plugin kind reserved for the built-in privacy component. pub const PII_REDACTION_PLUGIN_KIND: &str = "pii_redaction"; +/// Versioned inference contract implemented by PII detection providers. +pub const PII_DETECTION_PROVIDER_CONTRACT: &str = "nemo.relay.pii_detection.v1"; +pub(super) const DEFAULT_LOCAL_MODEL_LATENCY_MS: u64 = 250; +pub(super) const DEFAULT_LOCAL_MODEL_MIN_SCORE: f64 = 0.4; +pub(super) const MAX_LOCAL_MODEL_LATENCY_MS: u64 = 60_000; +pub(super) const MAX_LOCAL_MODEL_TARGET_PATHS: usize = 256; +pub(super) const MAX_LOCAL_MODEL_TARGET_PATH_BYTES: usize = 1024; +pub(super) const MAX_LOCAL_MODEL_REPLACEMENT_BYTES: usize = 1024; +pub(super) const MAX_LOCAL_MODEL_PROVIDER_VALUE_BYTES: usize = 1024; +pub(super) const MAX_LOCAL_MODEL_EXCLUDED_LABELS: usize = 128; +pub(super) const MAX_LOCAL_MODEL_LABEL_BYTES: usize = 128; +const BUILTIN_BACKEND_CONFIG_FIELDS: &[&str] = &[ + "preset", + "action", + "target_paths", + "pattern", + "detector", + "replacement", + "mask_char", + "unmasked_prefix", + "unmasked_suffix", + "custom_mark_payload_policy", +]; +const LOCAL_BACKEND_CONFIG_FIELDS: &[&str] = &[ + "backend", + "model_id", + "detector_profile", + "target_paths", + "target_path_patterns", + "min_score", + "excluded_labels", + "replacement", + "allow_network", + "max_latency_ms", +]; /// Top-level PII redaction component wrapper. #[derive(Debug, Clone)] @@ -237,23 +271,38 @@ impl Default for BuiltinBackendConfig { } } -/// Local-backend settings for a future in-process local-model runtime. +/// Local-backend settings for a same-machine local-model provider. #[derive(Debug, Clone, Default, Serialize, Deserialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] pub struct LocalBackendConfig { - /// Optional local-model backend identifier. + /// Registered local-model provider identifier. #[serde(default, skip_serializing_if = "Option::is_none")] pub backend: Option, - /// Optional model identifier reserved for future local-model runtimes. + /// Optional model identifier passed to the provider. #[serde(default, skip_serializing_if = "Option::is_none")] pub model_id: Option, - /// Optional detector profile reserved for future local-model runtimes. + /// Optional detector profile passed to the provider. #[serde(default, skip_serializing_if = "Option::is_none")] pub detector_profile: Option, - /// Whether a future local-model backend may use network calls. + /// Exact JSON-pointer paths to inspect. Empty means every string leaf. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub target_paths: Vec, + /// JSON-pointer patterns to inspect. A `*` segment matches one path segment. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub target_path_patterns: Vec, + /// Minimum provider confidence accepted for redaction. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub min_score: Option, + /// Provider labels that should not be redacted. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub excluded_labels: Vec, + /// Replacement applied to every accepted detection. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub replacement: Option, + /// Whether the provider may use network calls. #[serde(default, skip_serializing_if = "Option::is_none")] pub allow_network: Option, - /// Target latency budget hint for a future local-model backend. + /// Total provider deadline for one selected payload in milliseconds. #[serde(default, skip_serializing_if = "Option::is_none")] pub max_latency_ms: Option, } @@ -398,6 +447,11 @@ nemo_relay::editor_config! { backend => { label: "backend", kind: String, optional: true }, model_id => { label: "model_id", kind: String, optional: true }, detector_profile => { label: "detector_profile", kind: String, optional: true }, + target_paths => { label: "target_paths", kind: List, list: &nemo_relay::config_editor::STRING_LIST_ITEM }, + target_path_patterns => { label: "target_path_patterns", kind: List, list: &nemo_relay::config_editor::STRING_LIST_ITEM }, + min_score => { label: "min_score", kind: Float, optional: true }, + excluded_labels => { label: "excluded_labels", kind: List, list: &nemo_relay::config_editor::STRING_LIST_ITEM }, + replacement => { label: "replacement", kind: String, optional: true }, allow_network => { label: "allow_network", kind: Boolean, optional: true }, max_latency_ms => { label: "max_latency_ms", kind: Integer, optional: true }, } @@ -650,31 +704,14 @@ fn validate_pii_redaction_plugin_config_with_policy( &config.policy, plugin_config, "builtin", - &[ - "preset", - "action", - "target_paths", - "pattern", - "detector", - "replacement", - "mask_char", - "unmasked_prefix", - "unmasked_suffix", - "custom_mark_payload_policy", - ], + BUILTIN_BACKEND_CONFIG_FIELDS, ); validate_section_fields( &mut diagnostics, &config.policy, plugin_config, "local", - &[ - "backend", - "model_id", - "detector_profile", - "allow_network", - "max_latency_ms", - ], + LOCAL_BACKEND_CONFIG_FIELDS, ); validate_version(&mut diagnostics, &config.policy, config.version); validate_mode(&mut diagnostics, &config.policy, &config); @@ -683,6 +720,7 @@ fn validate_pii_redaction_plugin_config_with_policy( validate_builtin_mode_requirements(&mut diagnostics, &config.policy, plugin_config, &config); validate_builtin_action_requirements(&mut diagnostics, &config.policy, plugin_config, &config); validate_local_mode_requirements(&mut diagnostics, &config.policy, plugin_config, &config); + validate_local_backend_requirements(&mut diagnostics, &config.policy, &config); diagnostics } @@ -752,31 +790,14 @@ fn validate_profile_configuration( &config.policy, raw_profile, "builtin", - &[ - "preset", - "action", - "target_paths", - "pattern", - "detector", - "replacement", - "mask_char", - "unmasked_prefix", - "unmasked_suffix", - "custom_mark_payload_policy", - ], + BUILTIN_BACKEND_CONFIG_FIELDS, ); validate_section_fields( &mut profile_diagnostics, &config.policy, raw_profile, "local", - &[ - "backend", - "model_id", - "detector_profile", - "allow_network", - "max_latency_ms", - ], + LOCAL_BACKEND_CONFIG_FIELDS, ); validate_mode(&mut profile_diagnostics, &config.policy, &profile_config); validate_builtin_mode_requirements( @@ -797,16 +818,11 @@ fn validate_profile_configuration( raw_profile, &profile_config, ); - if profile.mode == "local_model" && !raw_profile.contains_key("local") { - push_policy_diag( - &mut profile_diagnostics, - config.policy.unsupported_value, - "pii_redaction.unsupported_value", - Some(PII_REDACTION_PLUGIN_KIND.to_string()), - Some("local".to_string()), - "`local` settings are required for a local-model profile".to_string(), - ); - } + validate_local_backend_requirements( + &mut profile_diagnostics, + &config.policy, + &profile_config, + ); prefix_profile_diagnostics(&mut profile_diagnostics, index); diagnostics.extend(profile_diagnostics); } @@ -867,6 +883,16 @@ fn validate_local_mode_requirements( config: &PiiRedactionConfig, ) { if config.mode == "local_model" { + if !plugin_config.contains_key("local") { + push_policy_diag( + diagnostics, + policy.unsupported_value, + "pii_redaction.unsupported_value", + Some(PII_REDACTION_PLUGIN_KIND.to_string()), + Some("local".to_string()), + "`local` settings are required when mode = 'local_model'".to_string(), + ); + } return; } if !plugin_config.contains_key("local") { @@ -883,6 +909,58 @@ fn validate_local_mode_requirements( ); } +fn validate_local_backend_requirements( + diagnostics: &mut Vec, + policy: &ConfigPolicy, + config: &PiiRedactionConfig, +) { + if config.mode != "local_model" { + return; + } + let Some(local) = config.local.as_ref() else { + return; + }; + for violation in validate_local_backend_config(local) { + push_policy_diag( + diagnostics, + policy.unsupported_value, + "pii_redaction.unsupported_value", + Some(PII_REDACTION_PLUGIN_KIND.to_string()), + Some(violation.field.to_string()), + violation.message, + ); + } + if local.target_paths.is_empty() && local.target_path_patterns.is_empty() { + diagnostics.push(ConfigDiagnostic { + level: DiagnosticLevel::Warning, + code: "pii_redaction.local_model_all_paths".to_string(), + component: Some(PII_REDACTION_PLUGIN_KIND.to_string()), + field: Some("local.target_paths".to_string()), + message: "local-model PII redaction has no target paths and will inspect every string leaf; configure explicit content paths to avoid classifying identifiers and metadata".to_string(), + }); + } +} + +pub(super) fn is_valid_json_pointer(path: &str) -> bool { + if path.is_empty() { + return true; + } + if !path.starts_with('/') { + return false; + } + let mut bytes = path.as_bytes().iter().copied(); + while let Some(byte) = bytes.next() { + if byte == b'~' && !matches!(bytes.next(), Some(b'0' | b'1')) { + return false; + } + } + true +} + +pub(super) fn is_valid_json_pointer_pattern(path: &str) -> bool { + is_valid_json_pointer(path) +} + fn validate_builtin_mode_requirements( diagnostics: &mut Vec, policy: &ConfigPolicy, diff --git a/crates/pii-redaction/src/local.rs b/crates/pii-redaction/src/local.rs index b1198c43b..d60c56972 100644 --- a/crates/pii-redaction/src/local.rs +++ b/crates/pii-redaction/src/local.rs @@ -1,115 +1,990 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::sync::{Arc, LazyLock, Mutex, MutexGuard}; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use nemo_relay::api::event::Event; +use nemo_relay::api::llm::LlmRequest; +use nemo_relay::api::runtime::{ + BuiltinLlmCodec, EventSanitizeFn, LlmCodecIdentity, LlmSanitizeRequestFn, + LlmSanitizeResponseFn, ToolSanitizeFn, +}; +use nemo_relay::codec::resolve::{ + ProviderSurface, detect_response_surface, request_codec as build_request_codec, + response_codec as build_response_codec, +}; +use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay::plugin::{ - PluginError, PluginRegistrationContext, Result as PluginResult, rollback_registrations, + InferenceProvider, PluginError, PluginRegistrationContext, Result as PluginResult, }; +use serde::de::DeserializeOwned; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value as Json}; -use super::component::PiiRedactionConfig; -use super::component::profile_registration_prefix; +use super::component::{ + DEFAULT_LOCAL_MODEL_LATENCY_MS, DEFAULT_LOCAL_MODEL_MIN_SCORE, LocalBackendConfig, + MAX_LOCAL_MODEL_EXCLUDED_LABELS, MAX_LOCAL_MODEL_LABEL_BYTES, MAX_LOCAL_MODEL_LATENCY_MS, + MAX_LOCAL_MODEL_PROVIDER_VALUE_BYTES, MAX_LOCAL_MODEL_REPLACEMENT_BYTES, + MAX_LOCAL_MODEL_TARGET_PATH_BYTES, MAX_LOCAL_MODEL_TARGET_PATHS, + PII_DETECTION_PROVIDER_CONTRACT, PiiRedactionConfig, is_valid_json_pointer, + is_valid_json_pointer_pattern, profile_registration_prefix, +}; +use super::overlay::BuiltinCodecName; -#[doc(hidden)] -pub type LocalBackendProvider = Arc< - dyn Fn(PiiRedactionConfig, &mut PluginRegistrationContext) -> PluginResult<()> + Send + Sync, ->; +const LOCAL_MODEL_CONTRACT_VERSION: u32 = 1; +const MAX_BATCH_ITEMS: usize = 64; +const MAX_BATCH_BYTES: usize = 64 * 1024; +const MAX_TEXT_BYTES: usize = 16 * 1024; +const MAX_TEXTS_PER_PAYLOAD: usize = 256; +const MAX_PAYLOAD_TEXT_BYTES: usize = 256 * 1024; +const MAX_DETECTIONS_PER_TEXT: usize = 128; -static LOCAL_BACKEND_PROVIDER: LazyLock>> = - LazyLock::new(|| Mutex::new(None)); +#[derive(Clone)] +struct CompiledLocalBackend { + provider_name: Arc, + provider: InferenceProvider, + model_id: Option, + detector_profile: Option, + target_paths: Arc>>, + target_path_patterns: Arc>, + min_score: f64, + excluded_labels: Arc>, + replacement: Arc, + timeout: Duration, + legacy_surface: Option, +} -fn local_backend_provider_guard() -> PluginResult>> -{ - LOCAL_BACKEND_PROVIDER.lock().map_err(|e| { - PluginError::Internal(format!( - "PII redaction local backend provider lock poisoned: {e}" - )) - }) +#[derive(Clone)] +struct JsonPointerPattern { + segments: Vec, } -#[doc(hidden)] -pub fn register_local_backend_provider(provider: LocalBackendProvider) -> PluginResult<()> { - let mut guard = local_backend_provider_guard()?; - *guard = Some(provider); - Ok(()) +impl JsonPointerPattern { + fn compile(pattern: String) -> Self { + Self { + segments: compile_json_pointer(pattern), + } + } + + fn matches(&self, path: &[String]) -> bool { + self.segments.len() == path.len() + && self + .segments + .iter() + .zip(path) + .all(|(pattern, segment)| pattern == "*" || pattern == segment) + } } -#[doc(hidden)] -pub fn clear_local_backend_provider() -> PluginResult<()> { - let mut guard = local_backend_provider_guard()?; - *guard = None; - Ok(()) +#[derive(Serialize)] +struct LocalModelRequest<'a> { + version: u32, + #[serde(skip_serializing_if = "Option::is_none")] + model_id: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + detector_profile: Option<&'a str>, + texts: Vec>, } -pub(super) fn register_local_backend( - config: PiiRedactionConfig, - ctx: &mut PluginRegistrationContext, - profile_name: Option<&str>, -) -> PluginResult<()> { - let provider = local_backend_provider_guard()?.clone(); +#[derive(Serialize)] +struct LocalModelText<'a> { + id: u32, + text: &'a str, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct LocalModelResponse { + version: u32, + detections: Vec, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct LocalModelDetection { + text_id: u32, + start_utf8: usize, + end_utf8: usize, + label: String, + score: f64, +} + +struct SelectedText { + text: String, + eligible: bool, +} + +enum EventField { + Data, + CategoryProfile, + Metadata, +} + +impl CompiledLocalBackend { + fn new( + config: LocalBackendConfig, + codec_name: Option, + ctx: &PluginRegistrationContext, + ) -> PluginResult { + if let Some(violation) = validate_local_backend_config(&config).into_iter().next() { + return Err(PluginError::InvalidConfig(violation.message)); + } + let provider_name = config + .backend + .as_deref() + .map(str::trim) + .expect("validated local backend has a provider name") + .to_string(); + let min_score = config.min_score.unwrap_or(DEFAULT_LOCAL_MODEL_MIN_SCORE); + let replacement = config + .replacement + .unwrap_or_else(|| "[REDACTED]".to_string()); + let max_latency_ms = config + .max_latency_ms + .unwrap_or(DEFAULT_LOCAL_MODEL_LATENCY_MS); + let surface = match codec_name.as_deref() { + Some(name) => Some(ProviderSurface::from_codec_name(name).ok_or_else(|| { + PluginError::InvalidConfig(format!("unsupported codec '{name}'")) + })?), + None => None, + }; + let provider = ctx + .inference_provider(&provider_name, PII_DETECTION_PROVIDER_CONTRACT) + .map_err(|error| { + PluginError::RegistrationFailed(format!( + "PII redaction inference provider '{provider_name}' is unavailable: {error}" + )) + })?; + Ok(Self { + provider_name: Arc::new(provider_name), + provider, + model_id: config.model_id.map(|value| value.trim().to_string()), + detector_profile: config + .detector_profile + .map(|value| value.trim().to_string()), + target_paths: Arc::new( + config + .target_paths + .into_iter() + .map(compile_json_pointer) + .collect(), + ), + target_path_patterns: Arc::new( + config + .target_path_patterns + .into_iter() + .map(JsonPointerPattern::compile) + .collect(), + ), + min_score, + excluded_labels: Arc::new( + config + .excluded_labels + .into_iter() + .map(|label| label.trim().to_string()) + .collect(), + ), + replacement: Arc::new(replacement), + timeout: Duration::from_millis(max_latency_ms), + legacy_surface: surface, + }) + } + + fn sanitize_json(&self, value: Json) -> Json { + self.sanitize_json_values(vec![value]) + .pop() + .expect("single-value sanitization returns one value") + } + + fn sanitize_json_values(&self, values: Vec) -> Vec { + self.sanitize_json_roots( + values + .into_iter() + .map(|value| (Vec::new(), value)) + .collect(), + ) + } + + fn sanitize_json_roots(&self, mut roots: Vec<(Vec, Json)>) -> Vec { + let mut texts = Vec::new(); + let mut total_bytes = 0; + let mut within_budget = true; + for (path, value) in &mut roots { + self.collect_strings( + value, + path, + &mut texts, + &mut total_bytes, + &mut within_budget, + ); + } + let sanitized = self.sanitize_texts(texts); + let mut index = 0; + for (path, value) in &mut roots { + self.replace_strings(value, path, &sanitized, &mut index); + } + roots.into_iter().map(|(_, value)| value).collect() + } + + fn collect_strings( + &self, + value: &Json, + path: &mut Vec, + texts: &mut Vec, + total_bytes: &mut usize, + within_budget: &mut bool, + ) { + match value { + Json::String(text) if self.matches_path(path) && *within_budget => { + if texts.len() >= MAX_TEXTS_PER_PAYLOAD { + *within_budget = false; + return; + } + if text.len() > MAX_TEXT_BYTES { + texts.push(SelectedText { + text: self.replacement.as_str().to_string(), + eligible: false, + }); + return; + } + let Some(next_total) = total_bytes.checked_add(text.len()) else { + *within_budget = false; + return; + }; + if next_total > MAX_PAYLOAD_TEXT_BYTES { + *within_budget = false; + return; + } + *total_bytes = next_total; + texts.push(SelectedText { + text: text.clone(), + eligible: true, + }); + } + Json::Array(items) => { + for (index, item) in items.iter().enumerate() { + path.push(index.to_string()); + self.collect_strings(item, path, texts, total_bytes, within_budget); + path.pop(); + } + } + Json::Object(fields) => { + for (key, value) in fields { + path.push(super::builtin::escape_json_pointer_segment(key)); + self.collect_strings(value, path, texts, total_bytes, within_budget); + path.pop(); + } + } + _ => {} + } + } + + fn replace_strings( + &self, + value: &mut Json, + path: &mut Vec, + sanitized: &[String], + index: &mut usize, + ) { + match value { + Json::String(text) if self.matches_path(path) => { + if let Some(replacement) = sanitized.get(*index) { + *text = replacement.clone(); + } else { + *text = self.replacement.as_str().to_string(); + } + *index += 1; + } + Json::Array(items) => { + for (item_index, item) in items.iter_mut().enumerate() { + path.push(item_index.to_string()); + self.replace_strings(item, path, sanitized, index); + path.pop(); + } + } + Json::Object(fields) => { + for (key, value) in fields { + path.push(super::builtin::escape_json_pointer_segment(key)); + self.replace_strings(value, path, sanitized, index); + path.pop(); + } + } + _ => {} + } + } + + fn matches_path(&self, path: &[String]) -> bool { + (self.target_paths.is_empty() && self.target_path_patterns.is_empty()) + || self.target_paths.contains(path) + || self + .target_path_patterns + .iter() + .any(|pattern| pattern.matches(path)) + } + + fn sanitize_texts(&self, mut texts: Vec) -> Vec { + let eligible = texts + .iter() + .enumerate() + .filter_map(|(index, text)| text.eligible.then_some(index)) + .collect::>(); + + let mut cursor = 0; + let deadline = Instant::now() + self.timeout; + while cursor < eligible.len() { + let start = cursor; + let mut batch_bytes = 0; + while cursor < eligible.len() && cursor - start < MAX_BATCH_ITEMS { + let next_bytes = texts[eligible[cursor]].text.len(); + if cursor > start && batch_bytes + next_bytes > MAX_BATCH_BYTES { + break; + } + batch_bytes += next_bytes; + cursor += 1; + } + let batch = &eligible[start..cursor]; + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + for index in &eligible[start..] { + texts[*index].text = self.replacement.as_str().to_string(); + } + break; + } + match self.sanitize_batch(&texts, batch, remaining) { + Ok(replacements) => { + for (index, replacement) in replacements { + texts[index].text = replacement; + } + } + Err(_) => { + log::warn!( + target: "nemo_relay.plugin", + event = "local_model_provider_failed", + plugin_kind = "pii_redaction", + provider = self.provider_name.as_str(), + batch_size = batch.len(), + reason = "provider_or_response"; + "PII local-model provider failed closed" + ); + for index in batch { + texts[*index].text = self.replacement.as_str().to_string(); + } + } + } + } + texts.into_iter().map(|selected| selected.text).collect() + } + + fn sanitize_batch( + &self, + texts: &[SelectedText], + batch: &[usize], + timeout: Duration, + ) -> PluginResult> { + let request = LocalModelRequest { + version: LOCAL_MODEL_CONTRACT_VERSION, + model_id: self.model_id.as_deref(), + detector_profile: self.detector_profile.as_deref(), + texts: batch + .iter() + .map(|index| LocalModelText { + id: u32::try_from(*index).expect("bounded text index fits u32"), + text: &texts[*index].text, + }) + .collect(), + }; + let request = serde_json::to_value(request)?; + let response = self.provider.invoke(request, timeout)?; + let response: LocalModelResponse = serde_json::from_value(response).map_err(|error| { + PluginError::RegistrationFailed(format!( + "local-model provider returned an invalid detection response: {error}" + )) + })?; + self.apply_response(texts, batch, response) + } + + fn apply_response( + &self, + texts: &[SelectedText], + batch: &[usize], + response: LocalModelResponse, + ) -> PluginResult> { + if response.version != LOCAL_MODEL_CONTRACT_VERSION { + return Err(PluginError::RegistrationFailed(format!( + "unsupported local-model response version {}", + response.version + ))); + } + if response.detections.len() > batch.len() * MAX_DETECTIONS_PER_TEXT { + return Err(PluginError::RegistrationFailed( + "local-model response exceeded the detection limit".into(), + )); + } + let allowed_ids = batch + .iter() + .map(|index| u32::try_from(*index).expect("bounded text index fits u32")) + .collect::>(); + let mut detections = HashMap::>::new(); + for detection in response.detections { + if !allowed_ids.contains(&detection.text_id) { + return Err(PluginError::RegistrationFailed(format!( + "local-model response referenced unknown text id {}", + detection.text_id + ))); + } + if detection.label.trim().is_empty() + || detection.label.len() > MAX_LOCAL_MODEL_LABEL_BYTES + { + return Err(PluginError::RegistrationFailed( + "local-model response contained an invalid detection label".into(), + )); + } + if !detection.score.is_finite() || !(0.0..=1.0).contains(&detection.score) { + return Err(PluginError::RegistrationFailed( + "local-model response contained an invalid detection score".into(), + )); + } + let text_detections = detections.entry(detection.text_id).or_default(); + if text_detections.len() >= MAX_DETECTIONS_PER_TEXT { + return Err(PluginError::RegistrationFailed(format!( + "local-model response exceeded the per-text detection limit of {MAX_DETECTIONS_PER_TEXT}" + ))); + } + text_detections.push(detection); + } + let mut replacements = Vec::new(); + for index in batch { + let id = u32::try_from(*index).expect("bounded text index fits u32"); + let Some(mut spans) = detections.remove(&id) else { + continue; + }; + spans.sort_by_key(|span| (span.start_utf8, span.end_utf8)); + let text = &texts[*index].text; + let mut previous_end = 0; + for span in &spans { + if span.start_utf8 >= span.end_utf8 + || span.end_utf8 > text.len() + || !text.is_char_boundary(span.start_utf8) + || !text.is_char_boundary(span.end_utf8) + || span.start_utf8 < previous_end + { + return Err(PluginError::RegistrationFailed( + "local-model response contained invalid or overlapping UTF-8 spans".into(), + )); + } + previous_end = span.end_utf8; + } + spans.retain(|detection| { + detection.score >= self.min_score + && !self.excluded_labels.contains(&detection.label) + }); + if spans.is_empty() { + continue; + } + let mut redacted = text.clone(); + for span in spans.iter().rev() { + redacted.replace_range(span.start_utf8..span.end_utf8, self.replacement.as_str()); + } + replacements.push((*index, redacted)); + } + Ok(replacements) + } + + fn sanitize_request_with_codec( + &self, + codec: &dyn LlmCodec, + request: &LlmRequest, + ) -> Option { + let annotated = codec.decode(request).ok()?; + let annotated = serde_json::to_value(annotated).ok()?; + let (headers, annotated) = + self.sanitize_request_parts(request.headers.clone(), annotated)?; + let annotated = serde_json::from_value(annotated).ok()?; + let mut encoded = codec.encode(&annotated, request).ok()?; + encoded.headers = headers; + Some(encoded) + } + + fn sanitize_raw_request(&self, mut request: LlmRequest) -> Option { + let headers = std::mem::take(&mut request.headers); + let content = std::mem::take(&mut request.content); + let (headers, content) = self.sanitize_request_parts(headers, content)?; + request.headers = headers; + request.content = content; + Some(request) + } + + fn sanitize_request_parts( + &self, + headers: Map, + content: Json, + ) -> Option<(Map, Json)> { + let mut values = self.sanitize_json_roots(vec![ + (vec!["headers".to_string()], Json::Object(headers)), + (Vec::new(), content), + ]); + let content = values.pop()?; + let Json::Object(headers) = values.pop()? else { + return None; + }; + Some((headers, content)) + } + + fn sanitize_response_with_codec( + &self, + codec: &dyn LlmResponseCodec, + surface: ProviderSurface, + payload: Json, + ) -> Option { + let codec_name = BuiltinCodecName::from_provider_surface(surface); + let annotated = codec.decode_response(&payload).ok()?; + let sanitized = sanitize_serializable(self, annotated).ok()?; + Some(codec_name.overlay_response_payload(payload, &sanitized)) + } + + fn selected_surface(&self, codec: &LlmCodecIdentity) -> Option { + match codec { + LlmCodecIdentity::None => self.legacy_surface, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) => { + Some(ProviderSurface::OpenAIChat) + } + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses) => { + Some(ProviderSurface::OpenAIResponses) + } + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) => { + Some(ProviderSurface::AnthropicMessages) + } + LlmCodecIdentity::Runtime(_) | LlmCodecIdentity::Opaque => None, + } + } + + fn uses_compatible_legacy_response_codec(&self, payload: &Json) -> bool { + self.legacy_surface + .is_some_and(|surface| detect_response_surface(payload) == Some(surface)) + } - let Some(provider) = provider else { + fn log_codec_failure(&self, direction: &'static str, codec: &LlmCodecIdentity, reason: &str) { + let codec_kind = match codec { + LlmCodecIdentity::None => "none", + LlmCodecIdentity::BuiltIn(_) => "builtin", + LlmCodecIdentity::Runtime(_) => "runtime", + LlmCodecIdentity::Opaque => "opaque", + }; log::warn!( target: "nemo_relay.plugin", - event = "plugin_resource_access_failed", + event = "local_model_codec_failed", plugin_kind = "pii_redaction", - profile = profile_name.unwrap_or("legacy"), - resource_kind = "local_model_backend", - permission = "execute", - reason = "provider_unavailable"; - "Plugin resource access validation failed" + provider = self.provider_name.as_str(), + direction, + codec_kind, + reason; + "PII local-model payload omitted after codec failure" ); - return Err(PluginError::RegistrationFailed( - "PII redaction local-model backend is unavailable in this runtime".to_string(), - )); - }; - log::info!( - target: "nemo_relay.plugin", - event = "plugin_resource_access_pending", - plugin_kind = "pii_redaction", - profile = profile_name.unwrap_or("legacy"), - resource_kind = "local_model_backend", - permission = "execute"; - "Plugin resource access validation started" - ); - let mut scoped_context = profile_name.map(|profile_name| { - PluginRegistrationContext::with_namespace( - ctx.qualify_name(&format!("{}/", profile_registration_prefix(profile_name))), + } +} + +fn compile_json_pointer(pointer: String) -> Vec { + pointer.strip_prefix('/').map_or_else(Vec::new, |path| { + path.split('/').map(str::to_string).collect() + }) +} + +pub(super) fn register_local_backend( + config: PiiRedactionConfig, + ctx: &mut PluginRegistrationContext, + profile_name: Option<&str>, +) -> PluginResult<()> { + let local = config.local.clone().ok_or_else(|| { + PluginError::InvalidConfig( + "local settings are required when mode = 'local_model'".to_string(), ) - }); - let provider_context = scoped_context.as_mut().unwrap_or(ctx); - match provider(config, provider_context) { - Ok(()) => { - if let Some(scoped_context) = scoped_context { - ctx.extend_registrations(scoped_context.into_registrations()); - } - log::info!( - target: "nemo_relay.plugin", - event = "plugin_resource_access_validated", - plugin_kind = "pii_redaction", - profile = profile_name.unwrap_or("legacy"), - resource_kind = "local_model_backend", - permission = "execute"; - "Plugin resource access validated" + })?; + let backend = CompiledLocalBackend::new(local, config.codec.clone(), ctx)?; + + if config.mark { + ctx.register_mark_sanitize_guardrail( + ®istration_name(profile_name, "mark"), + config.priority, + event_sanitize_callback(backend.clone(), None), + )?; + } + if config.tool_input { + ctx.register_tool_sanitize_request_guardrail( + ®istration_name(profile_name, "tool_input"), + config.priority, + tool_sanitize_callback(backend.clone()), + )?; + } + if config.tool_output { + ctx.register_tool_sanitize_response_guardrail( + ®istration_name(profile_name, "tool_output"), + config.priority, + tool_sanitize_callback(backend.clone()), + )?; + } + if config.input { + ctx.register_llm_sanitize_request_guardrail( + ®istration_name(profile_name, "input"), + config.priority, + llm_sanitize_request_callback(backend.clone()), + )?; + } + if config.input || config.tool_input { + ctx.register_scope_sanitize_start_guardrail( + ®istration_name( + profile_name, + if profile_name.is_some() { + "scope_start" + } else { + "input" + }, + ), + config.priority, + event_sanitize_callback(backend.clone(), Some((config.input, config.tool_input))), + )?; + } + if config.output { + ctx.register_llm_sanitize_response_guardrail( + ®istration_name(profile_name, "output"), + config.priority, + llm_sanitize_response_callback(backend.clone()), + )?; + } + if config.output || config.tool_output { + ctx.register_scope_sanitize_end_guardrail( + ®istration_name( + profile_name, + if profile_name.is_some() { + "scope_end" + } else { + "output" + }, + ), + config.priority, + event_sanitize_callback(backend, Some((config.output, config.tool_output))), + )?; + } + Ok(()) +} + +fn tool_sanitize_callback(backend: CompiledLocalBackend) -> ToolSanitizeFn { + Arc::new(move |_name, payload| backend.sanitize_json(payload)) +} + +fn event_sanitize_callback( + backend: CompiledLocalBackend, + scope_categories: Option<(bool, bool)>, +) -> EventSanitizeFn { + Arc::new(move |event, mut fields| { + if scope_categories.is_some_and(|(sanitize_llm, sanitize_tool)| { + matches!(event, Event::Scope(_)) + && event + .category() + .is_some_and(|category| match category.as_str() { + "llm" => !sanitize_llm, + "tool" => !sanitize_tool, + _ => false, + }) + }) { + return fields; + } + let specialized_scope = matches!(event, Event::Scope(_)) + && event + .category() + .is_some_and(|category| matches!(category.as_str(), "tool" | "llm")); + + let mut selected = Vec::with_capacity(3); + if !specialized_scope && let Some(data) = fields.data.take() { + selected.push((EventField::Data, data)); + } + if !specialized_scope + && let Some(profile) = fields.category_profile.take() + && let Ok(profile) = serde_json::to_value(profile) + { + selected.push((EventField::CategoryProfile, profile)); + } + if let Some(metadata) = fields.metadata.take() { + selected.push((EventField::Metadata, metadata)); + } + + let values = selected + .iter_mut() + .map(|(_, value)| std::mem::take(value)) + .collect(); + for ((field, _), value) in selected + .into_iter() + .zip(backend.sanitize_json_values(values)) + { + match field { + EventField::Data => fields.data = Some(value), + EventField::CategoryProfile => { + fields.category_profile = serde_json::from_value(value).ok(); + } + EventField::Metadata => fields.metadata = Some(value), + } + } + fields + }) +} + +fn llm_sanitize_request_callback(backend: CompiledLocalBackend) -> LlmSanitizeRequestFn { + Arc::new(move |mut request, context| { + if backend.target_paths.is_empty() && backend.target_path_patterns.is_empty() { + request.content = backend.sanitize_json(request.content); + return Some(request); + } + if matches!(context.codec(), LlmCodecIdentity::None) && backend.legacy_surface.is_none() { + return backend.sanitize_raw_request(request); + } + let resolved = context.resolve_codec(); + let fallback = if resolved.is_none() { + backend + .selected_surface(context.codec()) + .map(build_request_codec) + } else { + None + }; + let sanitized = resolved + .as_deref() + .or(fallback.as_deref()) + .and_then(|codec| backend.sanitize_request_with_codec(codec, &request)); + if sanitized.is_none() { + backend.log_codec_failure( + "request", + context.codec(), + "codec decode, sanitize, or encode failure", ); - Ok(()) - } - Err(error) => { - if let Some(scoped_context) = scoped_context { - let mut registrations = scoped_context.into_registrations(); - rollback_registrations(&mut registrations); - } - log::warn!( - target: "nemo_relay.plugin", - event = "plugin_resource_access_failed", - plugin_kind = "pii_redaction", - profile = profile_name.unwrap_or("legacy"), - resource_kind = "local_model_backend", - permission = "execute", - reason = "initialization_failed"; - "Plugin resource access validation failed" + } + sanitized + }) +} + +fn llm_sanitize_response_callback(backend: CompiledLocalBackend) -> LlmSanitizeResponseFn { + Arc::new(move |payload, context| { + if backend.target_paths.is_empty() && backend.target_path_patterns.is_empty() { + return Some(backend.sanitize_json(payload)); + } + if matches!(context.codec(), LlmCodecIdentity::None) && backend.legacy_surface.is_none() { + return Some(backend.sanitize_json(payload)); + } + if matches!(context.codec(), LlmCodecIdentity::None) + && !backend.uses_compatible_legacy_response_codec(&payload) + { + backend.log_codec_failure("response", context.codec(), "no compatible legacy codec"); + return None; + } + let surface = backend.selected_surface(context.codec()); + let resolved = context.resolve_codec(); + let fallback = if resolved.is_none() { + surface.map(build_response_codec) + } else { + None + }; + let sanitized = surface + .zip(resolved.as_deref().or(fallback.as_deref())) + .and_then(|(surface, codec)| { + backend.sanitize_response_with_codec(codec, surface, payload) + }); + if sanitized.is_none() { + backend.log_codec_failure( + "response", + context.codec(), + "codec decode, sanitize, or encode failure", + ); + } + sanitized + }) +} + +fn sanitize_serializable(backend: &CompiledLocalBackend, value: T) -> PluginResult +where + T: Serialize + DeserializeOwned, +{ + let value = serde_json::to_value(value)?; + serde_json::from_value(backend.sanitize_json(value)).map_err(PluginError::from) +} + +fn registration_name(profile_name: Option<&str>, callback_name: &str) -> String { + profile_name.map_or_else( + || callback_name.to_string(), + |profile_name| { + format!( + "{}/{callback_name}", + profile_registration_prefix(profile_name) + ) + }, + ) +} + +pub(super) struct LocalConfigViolation { + pub(super) field: &'static str, + pub(super) message: String, +} + +pub(super) fn validate_local_backend_config( + config: &LocalBackendConfig, +) -> Vec { + let mut violations = Vec::new(); + let mut push = |field, message| violations.push(LocalConfigViolation { field, message }); + + match config.backend.as_deref().map(str::trim) { + None | Some("") => push( + "local.backend", + "local.backend is required when mode = 'local_model'".into(), + ), + Some(backend) if backend.len() > MAX_LOCAL_MODEL_PROVIDER_VALUE_BYTES => push( + "local.backend", + format!( + "local.backend must not exceed {MAX_LOCAL_MODEL_PROVIDER_VALUE_BYTES} UTF-8 bytes" + ), + ), + Some(_) => {} + } + if config.allow_network == Some(true) { + push( + "local.allow_network", + "worker-backed local-model providers must not use network inference".into(), + ); + } + match config.max_latency_ms { + Some(0) => push( + "local.max_latency_ms", + "local.max_latency_ms must be greater than zero".into(), + ), + Some(latency) if latency > MAX_LOCAL_MODEL_LATENCY_MS => push( + "local.max_latency_ms", + format!("local.max_latency_ms must not exceed {MAX_LOCAL_MODEL_LATENCY_MS}"), + ), + _ => {} + } + for (field, value) in [ + ("local.model_id", config.model_id.as_deref()), + ("local.detector_profile", config.detector_profile.as_deref()), + ] { + if value.is_some_and(|value| value.trim().is_empty()) { + push(field, format!("{field} must be a non-empty string")); + } else if value.is_some_and(|value| value.len() > MAX_LOCAL_MODEL_PROVIDER_VALUE_BYTES) { + push( + field, + format!( + "{field} must not exceed {MAX_LOCAL_MODEL_PROVIDER_VALUE_BYTES} UTF-8 bytes" + ), ); - Err(error) } } + if config.target_paths.len() + config.target_path_patterns.len() > MAX_LOCAL_MODEL_TARGET_PATHS + { + push( + "local.target_paths", + format!( + "local.target_paths and local.target_path_patterns must contain at most {MAX_LOCAL_MODEL_TARGET_PATHS} entries in total" + ), + ); + } + if config + .target_paths + .iter() + .any(|path| path.len() > MAX_LOCAL_MODEL_TARGET_PATH_BYTES) + { + push( + "local.target_paths", + format!( + "local.target_paths entries must not exceed {MAX_LOCAL_MODEL_TARGET_PATH_BYTES} UTF-8 bytes" + ), + ); + } + if config + .target_paths + .iter() + .any(|path| !is_valid_json_pointer(path)) + { + push( + "local.target_paths", + "local.target_paths entries must be valid JSON pointers".into(), + ); + } + if config + .target_path_patterns + .iter() + .any(|path| path.len() > MAX_LOCAL_MODEL_TARGET_PATH_BYTES) + { + push( + "local.target_path_patterns", + format!( + "local.target_path_patterns entries must not exceed {MAX_LOCAL_MODEL_TARGET_PATH_BYTES} UTF-8 bytes" + ), + ); + } + if config + .target_path_patterns + .iter() + .any(|path| !is_valid_json_pointer_pattern(path)) + { + push( + "local.target_path_patterns", + "local.target_path_patterns entries must be valid JSON-pointer patterns".into(), + ); + } + if config + .min_score + .is_some_and(|score| !score.is_finite() || !(0.0..=1.0).contains(&score)) + { + push( + "local.min_score", + "local.min_score must be a finite number between 0 and 1".into(), + ); + } + if config.excluded_labels.len() > MAX_LOCAL_MODEL_EXCLUDED_LABELS { + push( + "local.excluded_labels", + format!( + "local.excluded_labels must contain at most {MAX_LOCAL_MODEL_EXCLUDED_LABELS} entries" + ), + ); + } + if config + .excluded_labels + .iter() + .any(|label| label.trim().is_empty() || label.len() > MAX_LOCAL_MODEL_LABEL_BYTES) + { + push( + "local.excluded_labels", + format!( + "local.excluded_labels entries must be non-empty and at most {MAX_LOCAL_MODEL_LABEL_BYTES} UTF-8 bytes" + ), + ); + } + if config + .excluded_labels + .iter() + .map(|label| label.trim()) + .collect::>() + .len() + != config.excluded_labels.len() + { + push( + "local.excluded_labels", + "local.excluded_labels must not contain duplicates".into(), + ); + } + if config + .replacement + .as_ref() + .is_some_and(|replacement| replacement.len() > MAX_LOCAL_MODEL_REPLACEMENT_BYTES) + { + push( + "local.replacement", + format!( + "local.replacement must not exceed {MAX_LOCAL_MODEL_REPLACEMENT_BYTES} UTF-8 bytes" + ), + ); + } + + violations } + +#[cfg(test)] +#[path = "../tests/unit/local_tests.rs"] +mod tests; diff --git a/crates/pii-redaction/tests/unit/component_tests.rs b/crates/pii-redaction/tests/unit/component_tests.rs index cd7ad4026..51d86a0d6 100644 --- a/crates/pii-redaction/tests/unit/component_tests.rs +++ b/crates/pii-redaction/tests/unit/component_tests.rs @@ -28,10 +28,12 @@ use crate::codec::openai_responses::OpenAIResponsesCodec; use crate::codec::request::AnnotatedLlmRequest; use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::plugin::{ - ConfigPolicy, DiagnosticLevel, PluginComponentSpec, PluginConfig, PluginError, + ConfigPolicy, DiagnosticLevel, InferenceProviderDescriptor, InferenceProviderRegistration, + InferenceProviderRegistry, PluginComponentSpec, PluginConfig, PluginError, PluginRegistrationContext, UnsupportedBehavior, clear_plugin_configuration, ensure_builtin_plugins_registered, initialize_plugins_exact as initialize_plugins, - list_plugin_kinds, rollback_registrations, validate_plugin_config, + initialize_plugins_with_inference_providers, list_plugin_kinds, rollback_registrations, + validate_plugin_config, }; use futures::StreamExt; use nemo_relay::observability::atif::{AtifAgentInfo, AtifExporter}; @@ -140,13 +142,32 @@ fn top_level_policy_controls_component_diagnostics() { fn reset_runtime() { enable_operational_logs(); let _ = clear_plugin_configuration(); - crate::plugins::pii_redaction::component::clear_local_backend_provider().unwrap(); crate::shared_runtime::reset_runtime_owner_for_tests(); let context = global_context(); *context.write().unwrap() = NemoRelayContextState::new(); register_pii_redaction_component().unwrap(); } +struct InferenceProviderGuard { + _registration: InferenceProviderRegistration, +} + +fn register_test_inference_provider( + registry: &InferenceProviderRegistry, + name: &str, + callback: impl Fn(Json, std::time::Duration) -> Result + Send + Sync + 'static, +) -> InferenceProviderGuard { + let registration = registry + .register( + InferenceProviderDescriptor::new(name, PII_DETECTION_PROVIDER_CONTRACT).unwrap(), + Arc::new(callback), + ) + .unwrap(); + InferenceProviderGuard { + _registration: registration, + } +} + fn setup_isolated_thread() { let stack = create_scope_stack(); set_thread_scope_stack(stack); @@ -296,6 +317,48 @@ impl LlmCodec for IdentifiedRequestCodec { } } +#[test] +fn raw_llm_paths_remain_usable_without_a_codec() { + let backend = crate::builtin::CompiledBuiltinBackend::new( + BuiltinBackendConfig { + action: "regex_replace".to_string(), + pattern: Some("sk-[A-Za-z0-9_-]+".to_string()), + replacement: Some("[REDACTED]".to_string()), + target_paths: vec!["/message".to_string()], + ..BuiltinBackendConfig::default() + }, + None, + ) + .unwrap(); + let sanitize_request = crate::builtin::llm_sanitize_request_callback(backend.clone()); + let sanitize_response = crate::builtin::llm_sanitize_response_callback(backend); + + let request = sanitize_request( + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "message": "sk-request-secret", + "model": "sk-model-identifier" + }), + }, + no_codec_request_context(), + ) + .expect("raw request paths should not require a codec"); + assert_eq!(request.content["message"], "[REDACTED]"); + assert_eq!(request.content["model"], "sk-model-identifier"); + + let response = sanitize_response( + json!({ + "message": "sk-response-secret", + "model": "sk-model-identifier" + }), + no_codec_context(), + ) + .expect("raw response paths should not require a codec"); + assert_eq!(response["message"], "[REDACTED]"); + assert_eq!(response["model"], "sk-model-identifier"); +} + #[test] fn normalized_llm_paths_use_the_active_codec_and_fail_closed_for_unknown_codecs() { let backend = crate::builtin::CompiledBuiltinBackend::new( @@ -1257,6 +1320,81 @@ fn profile_array_executes_every_profile_in_stable_array_order() { clear_plugin_configuration().unwrap(); } +#[test] +fn deterministic_and_local_model_profiles_compose_in_priority_order() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + let inference_providers = InferenceProviderRegistry::default(); + let _provider = + register_test_inference_provider(&inference_providers, "contextual", |request, _| { + assert_eq!( + request["texts"][0]["text"], "Alice emailed [REDACTED]", + "the local provider must receive the deterministic profile's output" + ); + Ok(json!({ + "version": 1, + "detections": [{ + "text_id": 0, + "start_utf8": 0, + "end_utf8": 5, + "label": "GIVEN_NAME", + "score": 0.99 + }] + })) + }); + + futures::executor::block_on(initialize_plugins_with_inference_providers( + plugin_config(json!({ + "codec": "openai_chat", + "profiles": [ + { + "mode": "builtin", + "priority": 80, + "builtin": { + "action": "redact", + "detector": "email" + } + }, + { + "mode": "local_model", + "priority": 90, + "local": { + "backend": "contextual", + "target_paths": ["/message"] + } + } + ] + })), + inference_providers, + )) + .unwrap(); + + let events = capture_events("pii-profile-composition"); + event( + EmitMarkEventParams::builder() + .name("composed-profile-mark") + .data(json!({ + "message": "Alice emailed alice@example.com", + "region": "us-west-2" + })) + .build(), + ) + .unwrap(); + let captured = captured_events_snapshot(&events); + assert_eq!( + captured[0].data().unwrap(), + &json!({ + "message": "[REDACTED] emailed [REDACTED]", + "region": "us-west-2" + }) + ); + + deregister_subscriber("pii-profile-composition").unwrap(); + clear_plugin_configuration().unwrap(); +} + #[test] fn profile_array_rejects_legacy_fields_and_reports_profile_paths() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); @@ -1364,10 +1502,13 @@ fn local_profile_registrations_receive_generated_namespaces() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); reset_runtime(); - register_local_backend_provider(Arc::new(|_, ctx| { - ctx.register_mark_sanitize_guardrail("shared", 100, Arc::new(|_, fields| fields)) - })) - .unwrap(); + let inference_providers = InferenceProviderRegistry::default(); + let _one = register_test_inference_provider(&inference_providers, "one", |_, _| { + Ok(json!({"version": 1, "detections": []})) + }); + let _two = register_test_inference_provider(&inference_providers, "two", |_, _| { + Ok(json!({"version": 1, "detections": []})) + }); let plugin = PiiRedactionPlugin; let config = json!({ @@ -1380,12 +1521,15 @@ fn local_profile_registrations_receive_generated_namespaces() { let Json::Object(config) = config else { panic!("component config must be object"); }; - let mut ctx = PluginRegistrationContext::with_namespace("profiles::"); + let mut ctx = PluginRegistrationContext::with_inference_providers( + Some("profiles::".into()), + inference_providers, + ); futures::executor::block_on(plugin.register(&config, &mut ctx)).unwrap(); let mut registrations = ctx.into_registrations(); let registrations_debug = format!("{registrations:?}"); - assert!(registrations_debug.contains("profiles::profile_00000000000000000000/shared")); - assert!(registrations_debug.contains("profiles::profile_00000000000000000001/shared")); + assert!(registrations_debug.contains("profiles::profile_00000000000000000000/mark")); + assert!(registrations_debug.contains("profiles::profile_00000000000000000001/mark")); rollback_registrations(&mut registrations); assert!(registrations.is_empty()); } @@ -1396,12 +1540,6 @@ fn failed_later_profile_rolls_back_earlier_profile_registrations() { reset_runtime(); setup_isolated_thread(); - register_local_backend_provider(Arc::new(|_, _| { - Err(PluginError::RegistrationFailed( - "intentional profile failure".into(), - )) - })) - .unwrap(); let activation = futures::executor::block_on(initialize_plugins(plugin_config(json!({ "codec": "openai_chat", "profiles": [ @@ -1696,6 +1834,112 @@ fn validate_rejects_local_section_outside_local_mode() { })); } +#[test] +fn validate_rejects_invalid_local_model_provider_settings() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + + let cases = [ + ( + json!({"mode": "local_model"}), + "local", + "required when mode = 'local_model'", + ), + ( + json!({"mode": "local_model", "local": {"backend": " "}}), + "local.backend", + "local.backend is required", + ), + ( + json!({ + "mode": "local_model", + "local": { + "backend": "x".repeat(MAX_LOCAL_MODEL_PROVIDER_VALUE_BYTES + 1) + } + }), + "local.backend", + "must not exceed", + ), + ( + json!({ + "mode": "local_model", + "local": {"backend": "worker", "allow_network": true} + }), + "local.allow_network", + "must not use network inference", + ), + ( + json!({ + "mode": "local_model", + "local": {"backend": "worker", "max_latency_ms": 0} + }), + "local.max_latency_ms", + "must be greater than zero", + ), + ( + json!({ + "mode": "local_model", + "local": {"backend": "worker", "max_latency_ms": 60001} + }), + "local.max_latency_ms", + "must not exceed 60000", + ), + ( + json!({ + "mode": "local_model", + "local": {"backend": "worker", "model_id": " "} + }), + "local.model_id", + "must be a non-empty string", + ), + ( + json!({ + "mode": "local_model", + "local": {"backend": "worker", "target_paths": ["message"]} + }), + "local.target_paths", + "valid JSON pointers", + ), + ( + json!({ + "mode": "local_model", + "local": { + "backend": "worker", + "replacement": "x".repeat(MAX_LOCAL_MODEL_REPLACEMENT_BYTES + 1) + } + }), + "local.replacement", + "must not exceed", + ), + ]; + + for (config, field, message) in cases { + let report = validate_plugin_config(&plugin_config(config)); + assert!(report.diagnostics.iter().any(|diagnostic| { + diagnostic.field.as_deref() == Some(field) && diagnostic.message.contains(message) + })); + } +} + +#[test] +fn validate_warns_when_local_model_inspects_every_string_leaf() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + + let report = validate_plugin_config(&plugin_config(json!({ + "mode": "local_model", + "codec": "openai_chat", + "local": {"backend": "worker"} + }))); + + assert!(!report.has_errors(), "{:?}", report.diagnostics); + assert!(report.diagnostics.iter().any(|diagnostic| { + diagnostic.level == DiagnosticLevel::Warning + && diagnostic.code == "pii_redaction.local_model_all_paths" + && diagnostic.field.as_deref() == Some("local.target_paths") + })); +} + #[test] fn validate_rejects_builtin_mode_without_builtin_section() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); @@ -1832,59 +2076,124 @@ fn local_backend_provider_is_invoked_for_local_model_mode() { let called = Arc::new(AtomicBool::new(false)); let called_inner = Arc::clone(&called); - register_local_backend_provider(Arc::new( - move |config, _ctx: &mut PluginRegistrationContext| { + let inference_providers = InferenceProviderRegistry::default(); + let _provider = register_test_inference_provider( + &inference_providers, + "test-provider", + move |request, _| { called_inner.store(true, Ordering::SeqCst); - assert_eq!(config.mode, "local_model"); - Ok(()) + assert_eq!(request["version"], 1); + Ok(json!({"version": 1, "detections": []})) }, + ); + setup_isolated_thread(); + futures::executor::block_on(initialize_plugins_with_inference_providers( + plugin_config(json!({ + "mode": "local_model", + "input": false, + "output": false, + "mark": false, + "tool_input": true, + "tool_output": false, + "local": {"backend": "test-provider"} + })), + inference_providers, )) .unwrap(); + tool_call( + ToolCallParams::builder() + .name("test") + .args(json!({"text": "hello"})) + .build(), + ) + .unwrap(); + assert!(called.load(Ordering::SeqCst)); + clear_plugin_configuration().unwrap(); +} + +#[test] +fn local_backend_reports_missing_and_failed_provider_initialization() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); let plugin = PiiRedactionPlugin; - let mut ctx = PluginRegistrationContext::with_namespace("test::"); let config = json!({ "mode": "local_model", + "input": false, + "output": false, + "mark": false, "tool_input": true, + "tool_output": false, + "local": {"backend": "missing"} }); let Json::Object(config) = config else { panic!("component config must be object"); }; + let mut ctx = PluginRegistrationContext::with_namespace("missing::"); + let missing = futures::executor::block_on(plugin.register(&config, &mut ctx)) + .expect_err("missing local provider should fail registration"); + assert!(missing.to_string().contains("unavailable")); - futures::executor::block_on(plugin.register(&config, &mut ctx)).unwrap(); - - assert!(called.load(Ordering::SeqCst)); + let inference_providers = InferenceProviderRegistry::default(); + let _failed = register_test_inference_provider(&inference_providers, "failed", |_, _| { + Err(PluginError::RegistrationFailed("provider failed".into())) + }); + let config = json!({ + "mode": "local_model", + "input": false, + "output": false, + "mark": false, + "tool_input": true, + "tool_output": false, + "local": {"backend": "failed"} + }); + let Json::Object(config) = config else { + panic!("component config must be object"); + }; + let mut ctx = PluginRegistrationContext::with_inference_providers( + Some("failed::".into()), + inference_providers, + ); + futures::executor::block_on(plugin.register(&config, &mut ctx)) + .expect("provider availability should be checked at registration"); + let mut registrations = ctx.into_registrations(); + rollback_registrations(&mut registrations); } #[test] -fn local_backend_reports_missing_and_failed_provider_initialization() { +fn local_backend_rejects_provider_with_incompatible_contract() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); reset_runtime(); + let inference_providers = InferenceProviderRegistry::default(); + let _registration = inference_providers + .register( + InferenceProviderDescriptor::new("embedding", "acme.embedding.v1").unwrap(), + Arc::new(|request, _| Ok(request)), + ) + .unwrap(); let plugin = PiiRedactionPlugin; - let config = json!({"mode": "local_model"}); - let Json::Object(config) = config else { + let Json::Object(config) = json!({ + "mode": "local_model", + "input": false, + "output": false, + "mark": false, + "tool_input": true, + "tool_output": false, + "local": {"backend": "embedding"} + }) else { panic!("component config must be object"); }; - let mut ctx = PluginRegistrationContext::with_namespace("missing::"); - let missing = futures::executor::block_on(plugin.register(&config, &mut ctx)) - .expect_err("missing local provider should fail registration"); - assert!(missing.to_string().contains("unavailable")); - - register_local_backend_provider(Arc::new(|_, _| { - Err(PluginError::RegistrationFailed( - "provider initialization failed".into(), - )) - })) - .unwrap(); - let mut ctx = PluginRegistrationContext::with_namespace("failed::"); - let failed = futures::executor::block_on(plugin.register(&config, &mut ctx)) - .expect_err("failed local provider should fail registration"); - assert!( - failed - .to_string() - .contains("provider initialization failed") + let mut ctx = PluginRegistrationContext::with_inference_providers( + Some("mismatch::".into()), + inference_providers, ); + + let error = futures::executor::block_on(plugin.register(&config, &mut ctx)) + .expect_err("PII must reject providers implementing another contract"); + + assert!(error.to_string().contains("acme.embedding.v1")); + assert!(error.to_string().contains(PII_DETECTION_PROVIDER_CONTRACT)); } #[test] diff --git a/crates/pii-redaction/tests/unit/local_tests.rs b/crates/pii-redaction/tests/unit/local_tests.rs new file mode 100644 index 000000000..4669b8422 --- /dev/null +++ b/crates/pii-redaction/tests/unit/local_tests.rs @@ -0,0 +1,960 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::atomic::{AtomicUsize, Ordering}; + +use nemo_relay::api::event::{BaseEvent, CategoryProfile, Event, EventSanitizeFields, MarkEvent}; +use nemo_relay::api::runtime::{LlmSanitizeRequestContext, LlmSanitizeResponseContext}; +use nemo_relay::codec::openai_responses::OpenAIResponsesCodec; +use nemo_relay::codec::request::AnnotatedLlmRequest; +use nemo_relay::codec::resolve::{ + ProviderSurface, request_codec as build_request_codec, response_codec as build_response_codec, +}; +use nemo_relay::codec::traits::LlmCodec; +use nemo_relay::plugin::{ + InferenceProviderDescriptor, InferenceProviderRegistration, InferenceProviderRegistry, + PluginRegistrationContext, +}; +use serde_json::json; + +use super::*; + +struct ProviderGuard { + _registration: InferenceProviderRegistration, +} + +struct IdentifiedRequestCodec { + identity: LlmCodecIdentity, +} + +impl LlmCodec for IdentifiedRequestCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + self.identity.clone() + } + + fn decode(&self, request: &LlmRequest) -> nemo_relay::error::Result { + OpenAIResponsesCodec.decode(request) + } + + fn encode( + &self, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> nemo_relay::error::Result { + OpenAIResponsesCodec.encode(annotated, original) + } +} + +fn provider_context( + name: &'static str, + callback: impl Fn(Json, Duration) -> PluginResult + Send + Sync + 'static, +) -> (ProviderGuard, PluginRegistrationContext) { + let registry = InferenceProviderRegistry::default(); + let registration = registry + .register( + InferenceProviderDescriptor::new(name, PII_DETECTION_PROVIDER_CONTRACT).unwrap(), + Arc::new(callback), + ) + .unwrap(); + ( + ProviderGuard { + _registration: registration, + }, + PluginRegistrationContext::with_inference_providers(None, registry), + ) +} + +fn backend( + name: &'static str, + callback: impl Fn(Json, Duration) -> PluginResult + Send + Sync + 'static, +) -> (ProviderGuard, CompiledLocalBackend) { + let (provider, ctx) = provider_context(name, callback); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some(name.into()), + ..LocalBackendConfig::default() + }, + None, + &ctx, + ) + .unwrap(); + (provider, backend) +} + +fn alice_detector(request: Json, _timeout: Duration) -> PluginResult { + let mut detections = Vec::new(); + for item in request["texts"] + .as_array() + .expect("provider request should contain texts") + { + let text_id = item["id"].as_u64().expect("text id should be an integer"); + let text = item["text"].as_str().expect("text should be a string"); + for (start, value) in text.match_indices("Alice") { + detections.push(json!({ + "text_id": text_id, + "start_utf8": start, + "end_utf8": start + value.len(), + "label": "given_name", + "score": 0.99 + })); + } + } + Ok(json!({"version": 1, "detections": detections})) +} + +#[test] +fn applies_non_overlapping_utf8_byte_spans() { + let (_provider, backend) = backend("local-test-utf8", |_, _| { + Ok(json!({ + "version": 1, + "detections": [ + { + "text_id": 0, + "start_utf8": 0, + "end_utf8": 5, + "label": "given_name", + "score": 0.99 + }, + { + "text_id": 0, + "start_utf8": 6, + "end_utf8": 12, + "label": "surname", + "score": 0.98 + } + ] + })) + }); + + assert_eq!( + backend.sanitize_json(json!({"text": "José Rivera"})), + json!({"text": "[REDACTED] [REDACTED]"}) + ); +} + +#[test] +fn malformed_or_overlapping_spans_fail_closed_for_the_batch() { + let (_provider, backend) = backend("local-test-overlap", |_, _| { + Ok(json!({ + "version": 1, + "detections": [ + { + "text_id": 0, + "start_utf8": 0, + "end_utf8": 4, + "label": "name", + "score": 0.9 + }, + { + "text_id": 0, + "start_utf8": 3, + "end_utf8": 6, + "label": "name", + "score": 0.9 + } + ] + })) + }); + + assert_eq!( + backend.sanitize_json(json!({"first": "secret", "second": "safe"})), + json!({"first": "[REDACTED]", "second": "[REDACTED]"}) + ); +} + +#[test] +fn provider_errors_fail_closed_without_changing_unselected_paths() { + let (_provider, ctx) = provider_context("local-test-failure", |_, _| { + Err(PluginError::RegistrationFailed("boom".into())) + }); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-failure".into()), + target_paths: vec!["/selected".into()], + replacement: Some("[PRIVATE]".into()), + ..LocalBackendConfig::default() + }, + None, + &ctx, + ) + .unwrap(); + + assert_eq!( + backend.sanitize_json(json!({ + "selected": "secret", + "unselected": "preserve" + })), + json!({ + "selected": "[PRIVATE]", + "unselected": "preserve" + }) + ); +} + +#[test] +fn batches_provider_requests_and_preserves_no_detection_values() { + let calls = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&calls); + let (_provider, backend) = backend("local-test-batching", move |request, _| { + observed.fetch_add(1, Ordering::SeqCst); + assert!(request["texts"].as_array().unwrap().len() <= MAX_BATCH_ITEMS); + Ok(json!({"version": 1, "detections": []})) + }); + let values = (0..(MAX_BATCH_ITEMS + 1)) + .map(|index| Json::String(format!("value-{index}"))) + .collect(); + + let sanitized = backend.sanitize_json(Json::Array(values)); + + assert_eq!(sanitized[0], "value-0"); + assert_eq!( + sanitized[MAX_BATCH_ITEMS], + format!("value-{MAX_BATCH_ITEMS}") + ); + assert_eq!(calls.load(Ordering::SeqCst), 2); +} + +#[test] +fn batches_multiple_event_roots_into_one_provider_request() { + let calls = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&calls); + let (_provider, backend) = backend("local-test-multi-root", move |request, _| { + observed.fetch_add(1, Ordering::SeqCst); + assert_eq!(request["texts"].as_array().unwrap().len(), 3); + Ok(json!({"version": 1, "detections": []})) + }); + + let sanitized = backend.sanitize_json_values(vec![ + json!({"message": "first"}), + json!({"name": "second"}), + json!({"trace": "third"}), + ]); + + assert_eq!( + sanitized, + vec![ + json!({"message": "first"}), + json!({"name": "second"}), + json!({"trace": "third"}), + ] + ); + assert_eq!(calls.load(Ordering::SeqCst), 1); +} + +#[test] +fn event_callback_batches_all_selected_fields_into_one_provider_request() { + let calls = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&calls); + let (_provider, backend) = backend("local-test-event-batching", move |request, _| { + observed.fetch_add(1, Ordering::SeqCst); + assert_eq!(request["texts"].as_array().unwrap().len(), 3); + Ok(json!({"version": 1, "detections": []})) + }); + let callback = event_sanitize_callback(backend, None); + let event = Event::Mark(MarkEvent::new( + BaseEvent::builder().name("mark").build(), + None, + None, + )); + + let sanitized = callback( + &event, + EventSanitizeFields { + data: Some(json!({"message": "first"})), + category_profile: Some(CategoryProfile::builder().subtype("second").build()), + metadata: Some(json!({"trace": "third"})), + }, + ); + + assert_eq!(sanitized.data.unwrap()["message"], "first"); + assert_eq!( + sanitized.category_profile.unwrap().subtype.as_deref(), + Some("second") + ); + assert_eq!(sanitized.metadata.unwrap()["trace"], "third"); + assert_eq!(calls.load(Ordering::SeqCst), 1); +} + +#[test] +fn latency_budget_applies_to_the_entire_payload() { + let calls = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&calls); + let (_provider, ctx) = provider_context("local-test-total-deadline", move |_, timeout| { + observed.fetch_add(1, Ordering::SeqCst); + std::thread::sleep(timeout + Duration::from_millis(5)); + Err(PluginError::RegistrationFailed("timed out".into())) + }); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-total-deadline".into()), + max_latency_ms: Some(10), + ..LocalBackendConfig::default() + }, + None, + &ctx, + ) + .unwrap(); + let values = (0..(MAX_BATCH_ITEMS + 1)) + .map(|index| Json::String(format!("value-{index}"))) + .collect(); + + let sanitized = backend.sanitize_json(Json::Array(values)); + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert!( + sanitized + .as_array() + .unwrap() + .iter() + .all(|value| value == "[REDACTED]") + ); +} + +#[test] +fn oversized_text_is_redacted_without_calling_the_provider() { + let calls = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&calls); + let (_provider, backend) = backend("local-test-oversized", move |_, _| { + observed.fetch_add(1, Ordering::SeqCst); + Ok(json!({"version": 1, "detections": []})) + }); + + assert_eq!( + backend.sanitize_json(Json::String("x".repeat(MAX_TEXT_BYTES + 1))), + Json::String("[REDACTED]".into()) + ); + assert_eq!(calls.load(Ordering::SeqCst), 0); +} + +#[test] +fn oversized_text_does_not_shift_later_provider_results() { + let (_provider, backend) = backend("local-test-oversized-middle", |_, _| { + Ok(json!({"version": 1, "detections": []})) + }); + + assert_eq!( + backend.sanitize_json(json!(["first", "x".repeat(MAX_TEXT_BYTES + 1), "third"])), + json!(["first", "[REDACTED]", "third"]) + ); +} + +#[test] +fn payload_count_limit_fails_closed_without_unbounded_provider_calls() { + let calls = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&calls); + let (_provider, backend) = backend("local-test-count-limit", move |_, _| { + observed.fetch_add(1, Ordering::SeqCst); + Ok(json!({"version": 1, "detections": []})) + }); + let values = (0..(MAX_TEXTS_PER_PAYLOAD + 1)) + .map(|index| Json::String(format!("value-{index}"))) + .collect(); + + let sanitized = backend.sanitize_json(Json::Array(values)); + + assert_eq!(sanitized[0], "value-0"); + assert_eq!(sanitized[MAX_TEXTS_PER_PAYLOAD], "[REDACTED]"); + assert_eq!( + calls.load(Ordering::SeqCst), + MAX_TEXTS_PER_PAYLOAD.div_ceil(MAX_BATCH_ITEMS) + ); +} + +#[test] +fn payload_byte_limit_fails_closed_after_the_bounded_prefix() { + let calls = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&calls); + let (_provider, backend) = backend("local-test-byte-limit", move |_, _| { + observed.fetch_add(1, Ordering::SeqCst); + Ok(json!({"version": 1, "detections": []})) + }); + let accepted = MAX_PAYLOAD_TEXT_BYTES / MAX_TEXT_BYTES; + let values = (0..=accepted) + .map(|_| Json::String("x".repeat(MAX_TEXT_BYTES))) + .collect(); + + let sanitized = backend.sanitize_json(Json::Array(values)); + + assert_eq!(sanitized[accepted - 1], "x".repeat(MAX_TEXT_BYTES)); + assert_eq!(sanitized[accepted], "[REDACTED]"); + assert_eq!( + calls.load(Ordering::SeqCst), + MAX_PAYLOAD_TEXT_BYTES.div_ceil(MAX_BATCH_BYTES) + ); +} + +#[test] +fn non_utf8_boundary_detection_fails_closed() { + let (_provider, backend) = backend("local-test-utf8-boundary", |_, _| { + Ok(json!({ + "version": 1, + "detections": [{ + "text_id": 0, + "start_utf8": 1, + "end_utf8": 2, + "label": "invalid", + "score": 1.0 + }] + })) + }); + + assert_eq!( + backend.sanitize_json(json!("é")), + Json::String("[REDACTED]".into()) + ); +} + +#[test] +fn local_policy_rejects_malformed_or_unbounded_values() { + let (_provider, ctx) = provider_context("local-test-policy-bounds", |request, _| Ok(request)); + + for (config, expected) in [ + ( + LocalBackendConfig { + backend: Some("x".repeat(MAX_LOCAL_MODEL_PROVIDER_VALUE_BYTES + 1)), + ..LocalBackendConfig::default() + }, + "local.backend", + ), + ( + LocalBackendConfig { + backend: Some("local-test-policy-bounds".into()), + target_paths: vec!["message".into()], + ..LocalBackendConfig::default() + }, + "valid JSON pointer", + ), + ( + LocalBackendConfig { + backend: Some("local-test-policy-bounds".into()), + target_paths: vec!["/bad~escape".into()], + ..LocalBackendConfig::default() + }, + "valid JSON pointer", + ), + ( + LocalBackendConfig { + backend: Some("local-test-policy-bounds".into()), + target_path_patterns: vec!["messages/*/content".into()], + ..LocalBackendConfig::default() + }, + "valid JSON-pointer pattern", + ), + ( + LocalBackendConfig { + backend: Some("local-test-policy-bounds".into()), + replacement: Some("x".repeat(MAX_LOCAL_MODEL_REPLACEMENT_BYTES + 1)), + ..LocalBackendConfig::default() + }, + "local.replacement", + ), + ( + LocalBackendConfig { + backend: Some("local-test-policy-bounds".into()), + model_id: Some("x".repeat(MAX_LOCAL_MODEL_PROVIDER_VALUE_BYTES + 1)), + ..LocalBackendConfig::default() + }, + "local.model_id", + ), + ( + LocalBackendConfig { + backend: Some("local-test-policy-bounds".into()), + min_score: Some(f64::NAN), + ..LocalBackendConfig::default() + }, + "local.min_score", + ), + ( + LocalBackendConfig { + backend: Some("local-test-policy-bounds".into()), + excluded_labels: vec!["NAME".into(), "NAME".into()], + ..LocalBackendConfig::default() + }, + "local.excluded_labels", + ), + ] { + let error = CompiledLocalBackend::new(config, None, &ctx) + .err() + .expect("invalid local policy should fail"); + assert!(error.to_string().contains(expected), "{error}"); + } +} + +#[test] +fn local_policy_accepts_root_and_escaped_json_pointers() { + assert!(is_valid_json_pointer("")); + assert!(is_valid_json_pointer("/nested/a~1b/~0value")); +} + +#[test] +fn target_path_patterns_match_one_segment_without_widening_exact_paths() { + let (_provider, ctx) = provider_context("local-test-path-patterns", alice_detector); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-path-patterns".into()), + target_paths: vec!["/exact".into()], + target_path_patterns: vec!["/messages/*/content".into()], + ..LocalBackendConfig::default() + }, + None, + &ctx, + ) + .unwrap(); + + assert_eq!( + backend.sanitize_json(json!({ + "exact": "Alice", + "messages": [ + {"content": "Alice", "name": "Alice"}, + {"content": "Alice"} + ], + "nested": {"messages": [{"content": "Alice"}]} + })), + json!({ + "exact": "[REDACTED]", + "messages": [ + {"content": "[REDACTED]", "name": "Alice"}, + {"content": "[REDACTED]"} + ], + "nested": {"messages": [{"content": "Alice"}]} + }) + ); +} + +#[test] +fn exact_paths_match_escaped_object_keys() { + let (_provider, ctx) = provider_context("local-test-escaped-path", alice_detector); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-escaped-path".into()), + target_paths: vec!["/a~1b/~0name".into()], + ..LocalBackendConfig::default() + }, + None, + &ctx, + ) + .unwrap(); + + assert_eq!( + backend.sanitize_json(json!({ + "a/b": {"~name": "Alice", "name": "Alice"}, + "a~1b": {"~name": "Alice"} + })), + json!({ + "a/b": {"~name": "[REDACTED]", "name": "Alice"}, + "a~1b": {"~name": "Alice"} + }) + ); +} + +#[test] +fn raw_request_paths_batch_headers_and_content_without_a_codec() { + let calls = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&calls); + let (_provider, ctx) = provider_context("local-test-raw-request", move |request, timeout| { + observed.fetch_add(1, Ordering::SeqCst); + alice_detector(request, timeout) + }); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-raw-request".into()), + target_paths: vec!["/headers/x-user".into(), "/message".into()], + ..LocalBackendConfig::default() + }, + None, + &ctx, + ) + .unwrap(); + + let sanitized = llm_sanitize_request_callback(backend)( + LlmRequest { + headers: Map::from_iter([("x-user".into(), json!("Alice"))]), + content: json!({"message": "Hello Alice", "model": "Alice-model"}), + }, + LlmSanitizeRequestContext::default(), + ) + .expect("raw request paths should not require a codec"); + + assert_eq!(sanitized.headers["x-user"], "[REDACTED]"); + assert_eq!(sanitized.content["message"], "Hello [REDACTED]"); + assert_eq!(sanitized.content["model"], "Alice-model"); + assert_eq!(calls.load(Ordering::SeqCst), 1); +} + +#[test] +fn raw_response_paths_work_without_a_codec() { + let (_provider, ctx) = provider_context("local-test-raw-response", alice_detector); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-raw-response".into()), + target_paths: vec!["/message".into()], + ..LocalBackendConfig::default() + }, + None, + &ctx, + ) + .unwrap(); + + let sanitized = llm_sanitize_response_callback(backend)( + json!({"message": "Hello Alice", "model": "Alice-model"}), + LlmSanitizeResponseContext::default(), + ) + .expect("raw response paths should not require a codec"); + + assert_eq!(sanitized["message"], "Hello [REDACTED]"); + assert_eq!(sanitized["model"], "Alice-model"); +} + +#[test] +fn request_codec_classifies_only_normalized_content_patterns() { + let (_provider, ctx) = provider_context("local-test-openai-request", alice_detector); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-openai-request".into()), + target_path_patterns: vec![ + "/messages/*/content".into(), + "/messages/*/content/*/text".into(), + ], + ..LocalBackendConfig::default() + }, + Some("openai_chat".into()), + &ctx, + ) + .unwrap(); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "Alice-model", + "trace_id": "Alice-trace", + "messages": [ + {"role": "system", "content": "Keep this policy"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "Email Alice"}, + {"type": "image_url", "image_url": {"url": "https://Alice.invalid"}} + ] + } + ] + }), + }; + let codec = build_request_codec(ProviderSurface::OpenAIChat); + let annotated = codec + .decode(&request) + .expect("OpenAI request should decode"); + let sanitized_annotated = + sanitize_serializable(&backend, annotated).expect("annotated request should sanitize"); + codec + .encode(&sanitized_annotated, &request) + .expect("sanitized OpenAI request should encode"); + + let sanitized = llm_sanitize_request_callback(backend)( + request, + LlmSanitizeRequestContext::for_request_codec(Some(codec)), + ) + .expect("valid request should remain observable"); + + assert_eq!(sanitized.content["model"], "Alice-model"); + assert_eq!(sanitized.content["trace_id"], "Alice-trace"); + assert_eq!( + sanitized.content["messages"][0]["content"], + "Keep this policy" + ); + assert_eq!( + sanitized.content["messages"][1]["content"][0]["text"], + "Email [REDACTED]" + ); + assert_eq!( + sanitized.content["messages"][1]["content"][1]["image_url"]["url"], + "https://Alice.invalid" + ); +} + +#[test] +fn request_uses_the_active_codec_instead_of_the_legacy_fallback() { + let (_provider, ctx) = provider_context("local-test-active-request-codec", alice_detector); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-active-request-codec".into()), + target_path_patterns: vec!["/messages/*/content".into()], + ..LocalBackendConfig::default() + }, + Some("openai_chat".into()), + &ctx, + ) + .unwrap(); + let sanitize = llm_sanitize_request_callback(backend); + let request = || LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "test-model", + "input": [{"role": "user", "content": "Email Alice"}] + }), + }; + + for codec in [ + Arc::new(IdentifiedRequestCodec { + identity: LlmCodecIdentity::Runtime("test.responses.v1".into()), + }) as Arc, + Arc::new(IdentifiedRequestCodec { + identity: LlmCodecIdentity::Opaque, + }) as Arc, + ] { + let sanitized = sanitize( + request(), + LlmSanitizeRequestContext::for_request_codec(Some(codec)), + ) + .expect("active runtime codecs should remain usable"); + assert_eq!(sanitized.content["input"][0]["content"], "Email [REDACTED]"); + } +} + +#[test] +fn response_codec_classifies_message_content_without_touching_identity_fields() { + let (_provider, ctx) = provider_context("local-test-openai-response", alice_detector); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-openai-response".into()), + target_path_patterns: vec!["/message".into(), "/message/*/text".into()], + ..LocalBackendConfig::default() + }, + Some("openai_chat".into()), + &ctx, + ) + .unwrap(); + let response = json!({ + "id": "Alice-response", + "model": "Alice-model", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Hello Alice"}, + "finish_reason": "stop" + }], + "vendor_trace": "Alice-trace" + }); + + let sanitized = backend + .sanitize_response_with_codec( + build_response_codec(ProviderSurface::OpenAIChat).as_ref(), + ProviderSurface::OpenAIChat, + response, + ) + .expect("configured codec should sanitize the response"); + + assert_eq!(sanitized["id"], "Alice-response"); + assert_eq!(sanitized["model"], "Alice-model"); + assert_eq!(sanitized["vendor_trace"], "Alice-trace"); + assert_eq!( + sanitized["choices"][0]["message"]["content"], + "Hello [REDACTED]" + ); +} + +#[test] +fn request_codec_failure_omits_the_observable_body() { + let (_provider, ctx) = provider_context("local-test-invalid-openai-request", alice_detector); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-invalid-openai-request".into()), + target_path_patterns: vec!["/messages/*/content".into()], + replacement: Some("[PRIVATE]".into()), + ..LocalBackendConfig::default() + }, + Some("openai_chat".into()), + &ctx, + ) + .unwrap(); + let request = LlmRequest { + headers: serde_json::Map::from_iter([( + "x-provider-id".into(), + Json::String("preserve-header".into()), + )]), + content: json!({ + "messages": "Alice cannot be decoded as an OpenAI message list", + "vendor_trace": "Alice-trace" + }), + }; + + let sanitized = llm_sanitize_request_callback(backend)( + request, + LlmSanitizeRequestContext::for_request_codec(Some(build_request_codec( + ProviderSurface::OpenAIChat, + ))), + ); + + assert!(sanitized.is_none()); +} + +#[test] +fn request_codec_ambiguous_multi_message_edit_fails_closed() { + let (_provider, ctx) = provider_context("local-test-ambiguous-openai-request", alice_detector); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-ambiguous-openai-request".into()), + target_path_patterns: vec!["/messages/*/content".into()], + replacement: Some("[PRIVATE]".into()), + ..LocalBackendConfig::default() + }, + Some("openai_chat".into()), + &ctx, + ) + .unwrap(); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "messages": [ + {"role": "system", "content": "Alice owns this policy"}, + {"role": "user", "content": "Email Alice"} + ] + }), + }; + + let sanitized = llm_sanitize_request_callback(backend)( + request, + LlmSanitizeRequestContext::for_request_codec(Some(build_request_codec( + ProviderSurface::OpenAIChat, + ))), + ); + + assert!(sanitized.is_none()); +} + +#[test] +fn response_codec_failure_omits_the_observable_payload() { + let (_provider, ctx) = provider_context("local-test-invalid-openai-response", alice_detector); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-invalid-openai-response".into()), + target_path_patterns: vec!["/message".into()], + replacement: Some("[PRIVATE]".into()), + ..LocalBackendConfig::default() + }, + Some("openai_chat".into()), + &ctx, + ) + .unwrap(); + let response = json!({ + "choices": "Alice cannot be decoded as an OpenAI response list", + "vendor_trace": "Alice-trace" + }); + + let sanitized = llm_sanitize_response_callback(backend)( + response, + LlmSanitizeResponseContext::for_response_codec(Some(build_response_codec( + ProviderSurface::OpenAIChat, + ))), + ); + + assert!(sanitized.is_none()); +} + +#[test] +fn host_policy_applies_score_threshold_and_label_exclusions() { + let (_provider, ctx) = provider_context("local-test-detection-policy", |_, _| { + Ok(json!({ + "version": 1, + "detections": [ + { + "text_id": 0, + "start_utf8": 0, + "end_utf8": 5, + "label": "LOW_SCORE", + "score": 0.49 + }, + { + "text_id": 0, + "start_utf8": 6, + "end_utf8": 11, + "label": "PRESERVE", + "score": 0.99 + }, + { + "text_id": 0, + "start_utf8": 12, + "end_utf8": 17, + "label": "REDACT", + "score": 0.99 + } + ] + })) + }); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-detection-policy".into()), + min_score: Some(0.5), + excluded_labels: vec!["PRESERVE".into()], + ..LocalBackendConfig::default() + }, + None, + &ctx, + ) + .unwrap(); + + assert_eq!( + backend.sanitize_json(json!("first next1 final")), + json!("first next1 [REDACTED]") + ); +} + +#[test] +fn validates_filtered_detections_before_applying_host_policy() { + let (_provider, ctx) = provider_context("local-test-filtered-invalid-span", |_, _| { + Ok(json!({ + "version": 1, + "detections": [{ + "text_id": 0, + "start_utf8": 0, + "end_utf8": 999, + "label": "LOW_SCORE", + "score": 0.1 + }] + })) + }); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-filtered-invalid-span".into()), + min_score: Some(0.5), + ..LocalBackendConfig::default() + }, + None, + &ctx, + ) + .unwrap(); + + assert_eq!( + backend.sanitize_json(json!("preserve without detections")), + json!("[REDACTED]") + ); +} + +#[test] +fn enforces_detection_limit_for_each_text() { + let (_provider, ctx) = provider_context("local-test-per-text-detection-limit", |_, _| { + let detections = (0..=MAX_DETECTIONS_PER_TEXT) + .map(|index| { + json!({ + "text_id": 0, + "start_utf8": index, + "end_utf8": index + 1, + "label": "NAME", + "score": 0.9 + }) + }) + .collect::>(); + Ok(json!({"version": 1, "detections": detections})) + }); + let backend = CompiledLocalBackend::new( + LocalBackendConfig { + backend: Some("local-test-per-text-detection-limit".into()), + ..LocalBackendConfig::default() + }, + None, + &ctx, + ) + .unwrap(); + + assert_eq!( + backend.sanitize_json(json!(["x".repeat(MAX_DETECTIONS_PER_TEXT + 1), "second"])), + json!(["[REDACTED]", "[REDACTED]"]) + ); +} diff --git a/crates/pii-redaction/tests/worker_provider_tests.rs b/crates/pii-redaction/tests/worker_provider_tests.rs new file mode 100644 index 000000000..3d6df4ee4 --- /dev/null +++ b/crates/pii-redaction/tests/worker_provider_tests.rs @@ -0,0 +1,369 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! End-to-end coverage for worker-backed local-model PII redaction. + +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::sync::{Arc, Mutex, OnceLock}; + +use nemo_relay::api::event::Event; +use nemo_relay::api::llm::{LlmCallExecuteParams, LlmRequest, llm_call_execute}; +use nemo_relay::api::runtime::LlmExecutionNextFn; +use nemo_relay::api::scope::{EmitMarkEventParams, event}; +use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; +use nemo_relay::codec::openai_chat::OpenAIChatCodec; +use nemo_relay::codec::traits::LlmResponseCodec; +use nemo_relay::plugin::dynamic::{ + DynamicPluginActivationSpec, DynamicPluginKind, PluginHostActivation, +}; +use nemo_relay::plugin::{PluginComponentSpec, PluginConfig, clear_plugin_configuration}; +use nemo_relay_pii_redaction::component::{ + PII_DETECTION_PROVIDER_CONTRACT, PII_REDACTION_PLUGIN_KIND, register_pii_redaction_component, +}; +use serde_json::{Map, json}; +use tempfile::TempDir; + +static WORKER_PII_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + +#[tokio::test(flavor = "multi_thread")] +async fn worker_provider_sanitizes_events_and_is_removed_after_host_clear() { + let _guard = WORKER_PII_TEST_LOCK.lock().await; + let _ = clear_plugin_configuration(); + register_pii_redaction_component().expect("PII component should register"); + let worker_binary = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_worker_manifest(&worker_binary); + + let plugin_config = PluginConfig { + version: 1, + components: vec![PluginComponentSpec { + kind: PII_REDACTION_PLUGIN_KIND.into(), + enabled: true, + config: Map::from_iter([ + ("mode".into(), json!("local_model")), + ("codec".into(), json!("openai_chat")), + ("input".into(), json!(true)), + ("output".into(), json!(true)), + ("tool_input".into(), json!(false)), + ("tool_output".into(), json!(false)), + ("mark".into(), json!(true)), + ( + "local".into(), + json!({ + "backend": "fixture_worker/fixture_local_model", + "target_paths": ["/message"], + "target_path_patterns": [ + "/messages/*/content", + "/message" + ] + }), + ), + ]), + }], + policy: Default::default(), + }; + let (activation, report) = PluginHostActivation::activate( + plugin_config, + [DynamicPluginActivationSpec { + plugin_id: "fixture_worker".into(), + kind: DynamicPluginKind::Worker, + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + environment_ref: None, + config: Map::from_iter([("provider_only".into(), json!(true))]), + }], + ) + .await + .expect("worker and PII component should activate together"); + assert!(!report.has_errors()); + let inference_providers = activation.inference_providers(); + assert!( + inference_providers + .resolve( + "fixture_worker/fixture_local_model", + PII_DETECTION_PROVIDER_CONTRACT, + ) + .is_ok() + ); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + let subscriber_name = "worker-pii-e2e"; + register_subscriber( + subscriber_name, + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .expect("test subscriber should register"); + + let callback_value = json!({ + "message": "keep PRIVATE hidden", + "unselected": "PRIVATE remains outside the configured path" + }); + event( + EmitMarkEventParams::builder() + .name("worker-pii") + .data(callback_value.clone()) + .build(), + ) + .expect("mark should emit"); + flush_subscribers().expect("sanitized event should flush"); + + { + let captured = events.lock().unwrap(); + assert_eq!(captured.len(), 1); + assert_eq!( + captured[0].data(), + Some(&json!({ + "message": "keep [REDACTED] hidden", + "unselected": "PRIVATE remains outside the configured path" + })) + ); + } + assert_eq!( + callback_value["message"], "keep PRIVATE hidden", + "sanitization must not mutate caller-owned JSON" + ); + + let callback_request = Arc::new(Mutex::new(None)); + let observed_request = Arc::clone(&callback_request); + let response = json!({ + "id": "chatcmpl-PRIVATE", + "model": "model-PRIVATE", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": "answer for PRIVATE" + }, + "finish_reason": "stop" + }] + }); + let callback_response = response.clone(); + let callback: LlmExecutionNextFn = Arc::new(move |request| { + *observed_request.lock().unwrap() = Some(request.clone()); + let response = callback_response.clone(); + Box::pin(async move { Ok(response) }) + }); + let request = LlmRequest { + headers: Map::new(), + content: json!({ + "model": "model-PRIVATE", + "messages": [ + {"role": "system", "content": "policy"}, + {"role": "user", "content": "question from PRIVATE"} + ], + "vendor_trace": "trace-PRIVATE" + }), + }; + let response_codec: Arc = Arc::new(OpenAIChatCodec); + + let returned = llm_call_execute( + LlmCallExecuteParams::builder() + .name("openai") + .request(request.clone()) + .func(callback) + .response_codec(response_codec) + .build(), + ) + .await + .expect("LLM callback should complete"); + flush_subscribers().expect("LLM events should flush"); + + assert_eq!( + returned, response, + "sanitize guardrails must not rewrite callback values" + ); + assert_eq!( + callback_request.lock().unwrap().as_ref(), + Some(&request), + "sanitize guardrails must not rewrite provider requests" + ); + let captured = events.lock().unwrap(); + assert_eq!(captured.len(), 3); + assert_eq!( + captured[1].input().unwrap()["content"]["messages"][1]["content"], + "question from [REDACTED]" + ); + assert_eq!( + captured[1].input().unwrap()["content"]["model"], + "model-PRIVATE", + "model identifiers are outside the selected content paths" + ); + assert_eq!( + captured[1].input().unwrap()["content"]["vendor_trace"], + "trace-PRIVATE", + "provider metadata is outside the selected content paths" + ); + assert_eq!( + captured[2].output().unwrap()["choices"][0]["message"]["content"], + "answer for [REDACTED]" + ); + assert_eq!(captured[2].output().unwrap()["id"], "chatcmpl-PRIVATE"); + assert_eq!(captured[2].output().unwrap()["model"], "model-PRIVATE"); + drop(captured); + + deregister_subscriber(subscriber_name).expect("test subscriber should deregister"); + activation.clear().expect("plugin host should clear"); + assert!( + inference_providers + .resolve( + "fixture_worker/fixture_local_model", + PII_DETECTION_PROVIDER_CONTRACT, + ) + .is_err(), + "worker provider should not outlive its host activation" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn worker_exit_during_sanitization_fails_closed_and_removes_provider() { + let _guard = WORKER_PII_TEST_LOCK.lock().await; + let _ = clear_plugin_configuration(); + register_pii_redaction_component().expect("PII component should register"); + let worker_binary = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_worker_manifest(&worker_binary); + let plugin_config = PluginConfig { + version: 1, + components: vec![PluginComponentSpec { + kind: PII_REDACTION_PLUGIN_KIND.into(), + enabled: true, + config: Map::from_iter([ + ("mode".into(), json!("local_model")), + ("input".into(), json!(false)), + ("output".into(), json!(false)), + ("tool_input".into(), json!(false)), + ("tool_output".into(), json!(false)), + ("mark".into(), json!(true)), + ( + "local".into(), + json!({ + "backend": "fixture_worker/fixture_local_model", + "target_paths": ["/message"] + }), + ), + ]), + }], + policy: Default::default(), + }; + let (activation, report) = PluginHostActivation::activate( + plugin_config, + [DynamicPluginActivationSpec { + plugin_id: "fixture_worker".into(), + kind: DynamicPluginKind::Worker, + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + environment_ref: None, + config: Map::from_iter([ + ("provider_only".into(), json!(true)), + ("exit_in_local_model".into(), json!(true)), + ]), + }], + ) + .await + .expect("worker and PII component should activate together"); + assert!(!report.has_errors()); + let inference_providers = activation.inference_providers(); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + let subscriber_name = "worker-pii-exit"; + register_subscriber( + subscriber_name, + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .expect("test subscriber should register"); + + event( + EmitMarkEventParams::builder() + .name("worker-pii-exit") + .data(json!({ + "message": "PRIVATE", + "unselected": "PRIVATE" + })) + .build(), + ) + .expect("worker failure must not block event emission"); + flush_subscribers().expect("fail-closed event should flush"); + assert_eq!( + events.lock().unwrap()[0].data(), + Some(&json!({ + "message": "[REDACTED]", + "unselected": "PRIVATE" + })) + ); + + deregister_subscriber(subscriber_name).expect("test subscriber should deregister"); + let error = activation + .clear() + .expect_err("stopped worker shutdown should be reported") + .to_string(); + assert!(error.contains("shutdown"), "{error}"); + assert!( + inference_providers + .resolve( + "fixture_worker/fixture_local_model", + PII_DETECTION_PROVIDER_CONTRACT, + ) + .is_err(), + "failed worker provider should not survive host teardown" + ); +} + +fn build_fixture_worker() -> PathBuf { + static FIXTURE_BINARY: OnceLock = OnceLock::new(); + FIXTURE_BINARY + .get_or_init(|| { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../core/tests/fixtures/worker_plugin/Cargo.toml"); + let target_dir = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../target/worker-plugin-fixture/target"); + let status = Command::new("cargo") + .arg("build") + .arg("--quiet") + .arg("--locked") + .arg("--manifest-path") + .arg(&manifest) + .arg("--target-dir") + .arg(&target_dir) + .status() + .expect("fixture worker build should start"); + assert!(status.success(), "fixture worker build should succeed"); + let binary = target_dir.join("debug").join(format!( + "nemo-relay-worker-plugin-fixture{}", + std::env::consts::EXE_SUFFIX + )); + assert!(binary.exists(), "fixture worker binary should exist"); + binary + }) + .clone() +} + +fn write_worker_manifest(binary: &Path) -> (TempDir, PathBuf) { + let temp = TempDir::new().expect("manifest directory should be created"); + let manifest = temp.path().join("relay-plugin.toml"); + let contents = format!( + r#" +manifest_version = 1 + +[plugin] +id = "fixture_worker" +kind = "worker" + +[compat] +relay = "={version}" +worker_protocol = "grpc-v1" + +[defaults] +enabled = false + +[capabilities] +items = ["plugin_worker"] + +[load] +runtime = "rust" +entrypoint = {entrypoint:?} +"#, + version = env!("CARGO_PKG_VERSION"), + entrypoint = binary.to_string_lossy(), + ); + std::fs::write(&manifest, contents).expect("worker manifest should be written"); + (temp, manifest) +} diff --git a/crates/worker-proto/README.md b/crates/worker-proto/README.md index d6e0e21d4..45c2aa3f2 100644 --- a/crates/worker-proto/README.md +++ b/crates/worker-proto/README.md @@ -33,6 +33,8 @@ tooling. from `v1` without generating protobuf code in a consumer project. - **Keep data ownership clear**: Carry Relay DTOs in JSON envelopes backed by `nemo-relay-types`; protobuf owns transport control flow. +- **Implement workers in another language**: Generate a `grpc-v1` service from + the protobuf when no maintained language-specific authoring SDK exists. ## What You Get @@ -41,6 +43,9 @@ tooling. clients, servers, services, and messages. - **JSON envelope helpers**: `json_envelope` and `decode_json_envelope` for serializing Relay DTOs into protocol payloads. +- **Language-neutral provider surface**: `INFERENCE_PROVIDER` carries a + versioned contract identifier plus component-owned request and response JSON + without coupling the host to the worker implementation language. ## Installation 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 75307e56d..59d8fb3f7 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 @@ -52,6 +52,7 @@ enum RegistrationSurface { MARK_SANITIZE_GUARDRAIL = 30; SCOPE_SANITIZE_START_GUARDRAIL = 31; SCOPE_SANITIZE_END_GUARDRAIL = 32; + INFERENCE_PROVIDER = 40; } enum LlmCodecKind { @@ -142,7 +143,7 @@ message Registration { RegistrationSurface surface = 2; int32 priority = 3; bool break_chain = 4; - reserved 5; + string contract = 5; } message InvokeRequest { @@ -158,6 +159,7 @@ message InvokeRequest { JsonEnvelope event = 10; ToolInvocation tool = 11; LlmInvocation llm = 12; + JsonEnvelope provider = 13; } } diff --git a/crates/worker-proto/tests/proto_tests.rs b/crates/worker-proto/tests/proto_tests.rs index 3b34ad9d1..6722c23ee 100644 --- a/crates/worker-proto/tests/proto_tests.rs +++ b/crates/worker-proto/tests/proto_tests.rs @@ -4,7 +4,8 @@ //! Tests for stable worker protocol helpers and enum values. use nemo_relay_worker_proto::v1::{ - HandshakeRequest, HealthRequest, InvokeRequest, JsonEnvelope, RegistrationSurface, ScopeType, + HandshakeRequest, HealthRequest, InvokeRequest, JsonEnvelope, Registration, + RegistrationSurface, ScopeType, }; use nemo_relay_worker_proto::{WORKER_PROTOCOL_GRPC_V1, decode_json_envelope, json_envelope}; use prost::Message; @@ -41,6 +42,13 @@ fn registration_surface_values_are_stable() { assert_eq!(RegistrationSurface::MarkSanitizeGuardrail as i32, 30); assert_eq!(RegistrationSurface::ScopeSanitizeStartGuardrail as i32, 31); assert_eq!(RegistrationSurface::ScopeSanitizeEndGuardrail as i32, 32); + assert_eq!(RegistrationSurface::InferenceProvider as i32, 40); + let encoded = Registration { + contract: "x".into(), + ..Default::default() + } + .encode_to_vec(); + assert_eq!(encoded, vec![42, 1, b'x']); } #[test] diff --git a/crates/worker/README.md b/crates/worker/README.md index 1f5d81ce3..16bb6fa9c 100644 --- a/crates/worker/README.md +++ b/crates/worker/README.md @@ -25,7 +25,8 @@ communicates with Relay through the versioned `grpc-v1` worker protocol. - **Isolate plugin code**: Run custom runtime behavior outside the Relay host process. - **Use typed registration APIs**: Implement `WorkerPlugin` and register - subscribers, guardrails, or intercepts with `PluginContext`. + subscribers, guardrails, intercepts, or inference providers with + `PluginContext`. - **Call the host runtime**: Emit marks, manage scopes, and invoke middleware continuations through `PluginRuntime`. - **Keep lifecycle managed**: Let Relay provide authenticated endpoints and @@ -85,6 +86,29 @@ Relay supplies the socket, activation ID, and authentication token through the worker environment. Use `serve_plugin` for Relay-spawned workers; explicit server configuration is intended for tests and custom launchers. +## Inference Providers + +A worker can expose detector or inference functionality to a first-party host +component without owning middleware policy: + +```rust +ctx.register_inference_provider( + "detector", + "acme.pii_detection.v1", + |request| async move { + Ok(serde_json::json!({ + "version": 1, + "detections": detect(request)? + })) + }, +); +``` + +The host publishes the provider as `/detector`. The consuming +component selects the exact contract and owns the request and response schema, +deadline, field selection, validation, and application of the result. The +worker callback should perform inference only. + ## Concurrency and Cancellation Unary and streaming callbacks run concurrently. Cancellation is cooperative: diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index f2e523f4a..900e0b5e9 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -270,6 +270,7 @@ type LlmRequestFn = Arc< type LlmExecutionFn = Arc BoxFutureResult + Send + Sync>; type LlmStreamExecutionFn = Arc BoxFutureResult + Send + Sync>; +type InferenceProviderFn = Arc BoxFutureResult + Send + Sync>; #[derive(Default)] struct WorkerHandlers { @@ -289,6 +290,7 @@ struct WorkerHandlers { llm_requests: HashMap, llm_executions: HashMap, llm_stream_executions: HashMap, + inference_providers: HashMap, } /// Registration context passed to [`WorkerPlugin::register`]. @@ -330,6 +332,23 @@ impl PluginContext { .insert(name.into(), Arc::new(callback)); } + /// Registers a named inference provider for a versioned host contract. + /// + /// The provider receives and returns versioned JSON data owned by the + /// consuming host component. It does not register middleware or decide + /// which runtime fields are sanitized. + pub fn register_inference_provider(&mut self, name: &str, contract: &str, callback: F) + where + F: Fn(Json) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + self.push_contract_registration(name, RegistrationSurface::InferenceProvider, contract); + self.handlers.inference_providers.insert( + name.into(), + Arc::new(move |request| Box::pin(callback(request))), + ); + } + fn register_event_sanitizer( &mut self, name: &str, @@ -666,6 +685,22 @@ impl PluginContext { surface: surface as i32, priority, break_chain, + contract: String::new(), + }); + } + + fn push_contract_registration( + &mut self, + name: &str, + surface: RegistrationSurface, + contract: &str, + ) { + self.handlers.registrations.push(Registration { + local_name: name.into(), + surface: surface as i32, + priority: 0, + break_chain: false, + contract: contract.into(), }); } } @@ -1661,6 +1696,12 @@ impl WorkerService { | RegistrationSurface::LlmExecutionIntercept => { self.invoke_llm_response(request, &scope, surface).await } + RegistrationSurface::InferenceProvider => { + let payload = provider_payload(request.payload)?; + let handler = self.inference_provider(&request.registration_name)?; + let future = with_thread_scope(&scope, || handler(payload)); + Ok(json_response(future.await?)) + } RegistrationSurface::LlmStreamExecutionIntercept | RegistrationSurface::Unspecified => { Err(WorkerSdkError::InvalidInput( "surface must use InvokeStream or is unspecified".into(), @@ -2056,6 +2097,18 @@ impl WorkerService { WorkerSdkError::InvalidInput(format!("llm execution '{name}' not registered")) }) } + + fn inference_provider(&self, name: &str) -> Result { + self.handlers + .lock() + .map_err(|err| WorkerSdkError::Callback(format!("handler lock poisoned: {err}")))? + .inference_providers + .get(name) + .cloned() + .ok_or_else(|| { + WorkerSdkError::InvalidInput(format!("inference provider '{name}' not registered")) + }) + } } struct ToolPayload { @@ -2172,6 +2225,19 @@ fn llm_payload( } } +fn provider_payload( + payload: Option, +) -> Result { + match payload { + Some(nemo_relay_worker_proto::v1::invoke_request::Payload::Provider(value)) => { + decode_json_envelope::(&value).map_err(Into::into) + } + _ => Err(WorkerSdkError::InvalidInput( + "expected inference provider payload".into(), + )), + } +} + fn required_json( value: Option, field: &str, @@ -2495,6 +2561,7 @@ fn all_surfaces() -> Vec { RegistrationSurface::MarkSanitizeGuardrail, RegistrationSurface::ScopeSanitizeStartGuardrail, RegistrationSurface::ScopeSanitizeEndGuardrail, + RegistrationSurface::InferenceProvider, ] } diff --git a/crates/worker/tests/worker_sdk_tests.rs b/crates/worker/tests/worker_sdk_tests.rs index 80653155f..462523b8b 100644 --- a/crates/worker/tests/worker_sdk_tests.rs +++ b/crates/worker/tests/worker_sdk_tests.rs @@ -114,6 +114,11 @@ async fn worker_service_enforces_auth_and_reports_registrations() { .supported_surfaces .contains(&(RegistrationSurface::LlmStreamExecutionIntercept as i32)) ); + assert!( + handshake + .supported_surfaces + .contains(&(RegistrationSurface::InferenceProvider as i32)) + ); let bad_health = client .health(Request::new(HealthRequest { @@ -200,7 +205,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(), 21); + assert_eq!(registrations.len(), 22); for local_name in [ "llm-sanitize-request", "llm-sanitize-response", @@ -272,6 +277,33 @@ async fn worker_service_enforces_auth_and_reports_registrations() { handle.abort(); } +#[tokio::test(flavor = "multi_thread")] +async fn worker_service_invokes_inference_provider() { + let (handle, mut client) = spawn_worker( + Arc::new(SurfacePlugin::default()), + "http://127.0.0.1:9".into(), + ) + .await; + let registrations = register_plugin(&mut client).await; + assert!(registrations.iter().any(|registration| { + registration.local_name == "local-model" + && registration.surface == RegistrationSurface::InferenceProvider as i32 + && registration.contract == "test.echo.v1" + })); + + let response = invoke_json( + &mut client, + provider_invoke("local-model", json!({"text": "private"})), + ) + .await; + + assert_eq!( + response, + json!({"text": "private", "provider": "local-model"}) + ); + handle.abort(); +} + #[tokio::test(flavor = "multi_thread")] async fn worker_service_rejects_duplicate_registration_names_on_one_surface() { let (handle, mut client) = spawn_worker( @@ -1226,6 +1258,10 @@ async fn worker_service_reports_missing_handlers_and_malformed_payloads() { ), "llm execution", ), + ( + provider_invoke("missing-inference-provider", json!({})), + "inference provider", + ), ] { assert_worker_error( client @@ -1929,6 +1965,9 @@ impl WorkerPlugin for SurfacePlugin { ctx.register_llm_stream_execution_intercept("llm-stream-open-error", 1, |_, _, _| async { Err(WorkerSdkError::Callback("stream open boom".into())) }); + ctx.register_inference_provider("local-model", "test.echo.v1", |request| async move { + Ok(set_json_field(request, "provider", "local-model")) + }); Ok(()) } } @@ -2521,6 +2560,21 @@ fn tool_invoke( } } +fn provider_invoke(registration_name: &str, value: Json) -> InvokeRequest { + InvokeRequest { + activation_id: ACTIVATION_ID.into(), + invocation_id: "invoke-1".into(), + registration_name: registration_name.into(), + surface: RegistrationSurface::InferenceProvider as i32, + continuation_id: String::new(), + scope: Some(scope_context()), + auth_token: AUTH_TOKEN.into(), + payload: Some( + nemo_relay_worker_proto::v1::invoke_request::Payload::Provider(json_env(value)), + ), + } +} + fn llm_invoke( registration_name: &str, surface: RegistrationSurface, diff --git a/docs/about-nemo-relay/release-notes/index.mdx b/docs/about-nemo-relay/release-notes/index.mdx index a87eb2176..55663a6a1 100644 --- a/docs/about-nemo-relay/release-notes/index.mdx +++ b/docs/about-nemo-relay/release-notes/index.mdx @@ -96,8 +96,10 @@ For the new callback contract and codec operations, refer to - The NeMo Guardrails remote backend inherits its configured service's availability, latency, and policy behavior. The local backend requires Python 3.11 or later and `nemoguardrails==0.22.0`. -- The PII redaction plugin currently supports its deterministic local backend; - local-model backend configuration is reserved for future work. +- The PII redaction plugin supports deterministic built-in policies and + worker-backed local-model providers. Local-model mode requires a separately + installed compatible `grpc-v1` provider; the optional Rampart provider + supports Latin-script text and does not provide complete PII coverage. - Pricing and optimization estimates depend on model names, token data, pricing sources, and freshness evidence. Missing or inconsistent evidence produces partial or absent cost fields rather than zero values. diff --git a/docs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdx b/docs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdx index b8ca9bfbb..85a703139 100644 --- a/docs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdx +++ b/docs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdx @@ -22,8 +22,8 @@ Workers implement the `PluginWorker` service: - `Handshake` and `Health` identify a ready worker. - `Validate` returns configuration diagnostics. -- `Register` returns declarative subscriber, guardrail, and intercept - registrations. +- `Register` returns declarative subscriber, guardrail, intercept, and + inference-provider registrations. - `Invoke` and `InvokeStream` run registered behavior. - `CancelInvocation` requests cancellation, and `Shutdown` requests process termination. @@ -67,12 +67,19 @@ return `WorkerError` without registrations. The supported surfaces are: - `LLM_SANITIZE_REQUEST_GUARDRAIL`, `LLM_SANITIZE_RESPONSE_GUARDRAIL`, `LLM_CONDITIONAL_EXECUTION_GUARDRAIL`, `LLM_REQUEST_INTERCEPT`, `LLM_EXECUTION_INTERCEPT`, and `LLM_STREAM_EXECUTION_INTERCEPT` +- `MARK_SANITIZE_GUARDRAIL`, `SCOPE_SANITIZE_START_GUARDRAIL`, and + `SCOPE_SANITIZE_END_GUARDRAIL` +- `INFERENCE_PROVIDER` `InvokeRequest` identifies the registration, surface, invocation, optional continuation, and scope context. Its payload is one of an event, tool -invocation, or LLM invocation. `InvokeResponse` returns an empty result, JSON -result, guardrail result, LLM request-intercept result, tool-execution result, -or `WorkerError`. `InvokeStream` emits JSON chunks or `WorkerError` chunks. +invocation, LLM invocation, or component-owned provider request. +`InvokeResponse` returns an empty result, JSON result, guardrail result, LLM +request-intercept result, tool-execution result, or `WorkerError`. +`INFERENCE_PROVIDER` declares a versioned contract and uses a JSON request and +JSON result. The consuming first-party component selects that exact contract, +owns its schema, and resolves the provider as `/`. +`InvokeStream` emits JSON chunks or `WorkerError` chunks. Every LLM sanitizer invocation includes a directional context with tagged codec identity: `none`, `builtin(id)`, `runtime(id)`, or `opaque`. Worker SDKs expose diff --git a/docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx b/docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx index 79d20f405..3c7451108 100644 --- a/docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx +++ b/docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx @@ -165,6 +165,31 @@ registers a tool request intercept that updates the request JSON and emits a mark through `PluginContext.runtime`. The `main` entrypoint blocks until Relay requests shutdown. +### Provide local inference + +Use `register_inference_provider` when a first-party Relay component owns a +versioned contract and needs an isolated detector or inference implementation: + +```python +async def detect(request: Json) -> Json: + return { + "version": 1, + "detections": await model.detect(request["texts"]), + } + + +ctx.register_inference_provider( + "detector", + "acme.pii_detection.v1", + detect, +) +``` + +Relay publishes this registration as `/detector`. The consuming +component selects the exact contract and owns its JSON schema, field selection, +deadline, response validation, and application. The worker should perform +inference only. + ## Create the Manifest Create `relay-plugin.toml` with the following content: diff --git a/docs/build-plugins/dynamic-plugins/grpc-worker/rust/about.mdx b/docs/build-plugins/dynamic-plugins/grpc-worker/rust/about.mdx index 9c5fa114f..cf687a433 100644 --- a/docs/build-plugins/dynamic-plugins/grpc-worker/rust/about.mdx +++ b/docs/build-plugins/dynamic-plugins/grpc-worker/rust/about.mdx @@ -68,6 +68,29 @@ identity from `WorkerPlugin::plugin_id()`; custom launchers can use credentials and runs until Relay requests shutdown. Do not start the worker directly for normal operation. +### Provide local inference + +Use `register_inference_provider` when a first-party Relay component owns a +versioned contract and needs isolated inference: + +```rust +context.register_inference_provider( + "detector", + "acme.pii_detection.v1", + |request| async move { + Ok(serde_json::json!({ + "version": 1, + "detections": detect(request)? + })) + }, +); +``` + +Relay publishes this registration as `/detector`. The consuming +component selects the exact contract and owns its JSON schema, field selection, +deadline, response validation, and application. The worker should perform +inference only. + ## Package the Worker Create `relay-plugin.toml` with the following content. The artifact digest must diff --git a/docs/configure-plugins/pii-redaction/about.mdx b/docs/configure-plugins/pii-redaction/about.mdx index a37cd4812..3c6bd7f65 100644 --- a/docs/configure-plugins/pii-redaction/about.mdx +++ b/docs/configure-plugins/pii-redaction/about.mdx @@ -24,7 +24,8 @@ The plugin supports these backend modes: - `builtin` - Uses a native Rust backend for deterministic payload sanitization. - `local_model` - - Reserves a future local-model backend lane for more stochastic detection behavior. + - Delegates bounded detector inference to a manifest-backed `grpc-v1` + worker while the PII component retains sanitization policy. ## Use This Plugin When @@ -39,6 +40,8 @@ Start here when you need to: first-party NeMo Relay components. - Use built-in detector presets for common values such as emails, phone numbers, URLs, API keys, and IP addresses without writing custom regexes. +- Compose deterministic recognizers with an optional local contextual + classifier without loading model dependencies into the Relay host. ## Plugin Versus Middleware @@ -90,18 +93,21 @@ The current built-in backend supports five actions: The current backend boundary is intentional: - Managed tool surfaces are sanitized as JSON payloads with exact JSON-pointer - targeting. + targeting or bounded single-segment wildcard patterns. - Managed LLM requests use the active resolved codec for each call, including built-in, runtime, and opaque codecs. Normalized response projection requires a recognized built-in codec because response codec capabilities are decode-only. This lets redaction target normalized Relay shapes such as - `/messages/0/content` and `/message`. + `/messages/*/content` and `/message`. - `mark` sanitizes `data`, `category_profile`, and `metadata` independently on every mark event. It defaults to `true`; set `mark = false` to opt out. - `input`, `output`, `tool_input`, and `tool_output` sanitize scope metadata on their corresponding lifecycle events. Tool and LLM primary data and typed profiles remain on their specialized sanitizer paths to avoid applying the same mask or hash twice. +- Local-model workers receive only the selected string values. Relay owns + batching, deadlines, confidence and label policy, response validation, + replacement, and fail-closed behavior. ## Observability Boundary @@ -116,20 +122,25 @@ That means: For managed LLM requests, codec decode and re-encode can canonicalize the emitted provider-shaped start event. For example, Relay can record an OpenAI Responses request in the codec's canonical `input` array form rather than the -original shorthand form. +original shorthand form. If codec processing fails, the local-model backend +omits the emitted LLM body rather than applying normalized selectors to an +incompatible raw provider shape. ## Current Boundaries This plugin is intentionally scoped to a deterministic built-in backend plus a -future local-model extension point. +local worker-provider extension point. In particular: -- `local_model` is an extension point, not a complete backend implementation - today. +- Local-model providers are optional dynamic plugins. They register a + language-neutral JSON request-response contract through the `grpc-v1` + worker protocol. - The plugin does not mutate the real callback arguments or return values. -- The plugin does not add a subtree or prefix selector language beyond exact - JSON-pointer matching. +- `target_path_patterns` supports `*` as one complete JSON-pointer path + segment. It does not support recursive or partial-segment matching. +- `allow_network = false` prohibits network inference by contract, but the + worker process is not a network sandbox. Install only trusted providers. ## Pages diff --git a/docs/configure-plugins/pii-redaction/configuration.mdx b/docs/configure-plugins/pii-redaction/configuration.mdx index d10f84813..5bb3e4a31 100644 --- a/docs/configure-plugins/pii-redaction/configuration.mdx +++ b/docs/configure-plugins/pii-redaction/configuration.mdx @@ -128,14 +128,14 @@ The following table compares the available PII redaction backends: | Area | `builtin` | `local_model` | | --- | --- | --- | | Built-in component kind and config validation | Supported | Supported | -| Managed LLM `input` | Supported | Not implemented | -| Managed LLM `output` | Supported | Not implemented | -| Mark and generic scope event fields | Supported | Not implemented | -| Managed `tool_input` | Supported | Not implemented | -| Managed `tool_output` | Supported | Not implemented | +| Managed LLM `input` | Supported | Supported | +| Managed LLM `output` | Supported | Supported | +| Mark and generic scope event fields | Supported | Supported | +| Managed `tool_input` | Supported | Supported | +| Managed `tool_output` | Supported | Supported | | Built-in actions | `remove`, `redact`, `regex_replace`, `hash`, `mask` | N/A | -| Codec support | `openai_chat`, `openai_responses`, `anthropic_messages` | Runtime-specific future implementation | -| Runtime availability | Any runtime that includes the `nemo-relay-pii-redaction` plugin crate | Runtimes that install a local backend provider | +| Codec support | `openai_chat`, `openai_responses`, `anthropic_messages` | `openai_chat`, `openai_responses`, `anthropic_messages` | +| Runtime availability | Any runtime that includes the `nemo-relay-pii-redaction` plugin crate | Runtimes with an active `grpc-v1` worker provider | ## Built-in Mode @@ -217,7 +217,8 @@ Use the editor when you want to: - Set the LLM `codec` - Edit `builtin` action settings such as `action`, `target_paths`, `pattern`, `detector`, `replacement`, and masking fields -- Edit `local.backend` for a runtime-provided future local-model backend +- Edit worker-backed local-model settings such as `local.backend`, + `target_paths`, `target_path_patterns`, `min_score`, and `excluded_labels` The editor preserves unknown fields when it rewrites an existing `pii_redaction` component, so future or runtime-specific settings are not @@ -314,9 +315,120 @@ When `detector` is set and you do not specify `unmasked_prefix` or - `gcp_api_key`: Preserves the `AIza`-style prefix and the last four characters - `azure_storage_account_key`: Preserves the last four characters +## Local-Model Mode + +Use `local_model` when a manifest-backed `grpc-v1` worker should detect +contextual PII. The worker performs inference only. The PII component selects +fields, batches text, enforces deadlines, validates detections, applies +confidence and label policy, replaces accepted spans, and fails closed. + +The provider name is `/`. For example, a worker +with plugin ID `acme.pii_worker` that registers `detector` for +`nemo.relay.pii_detection.v1` is selected as `acme.pii_worker/detector`. Relay +rejects providers that declare another contract before installing the +sanitizer. + +```toml +[[components]] +kind = "pii_redaction" +enabled = true + +[components.config] +mode = "local_model" +codec = "openai_chat" + +[components.config.local] +backend = "acme.pii_worker/detector" +model_id = "acme-pii-v1" +detector_profile = "default" +min_score = 0.4 +excluded_labels = ["ORGANIZATION"] +target_path_patterns = [ + "/messages/*/content", + "/messages/*/content/*/text", + "/message", + "/message/*/text", +] +replacement = "[REDACTED]" +allow_network = false +max_latency_ms = 250 +``` + +Worker providers are installed before built-in components initialize and are +removed after their dependent sanitizers. `allow_network = true` is rejected: +this lane is for same-machine inference. This is a configuration contract, not +a process sandbox. + +Provider failures, timeouts, malformed responses, invalid UTF-8 spans, +overlapping spans, and input-limit violations fail closed for the affected +batch. If a configured codec cannot decode or safely re-encode an LLM payload, +Relay omits that request or response payload from the emitted event. The +default deadline is 250 ms for the complete selected payload, including every +provider batch. Configuration above 60 seconds is rejected. + +Use `profiles` to run deterministic recognizers before a contextual model: + +```toml +[components.config] +codec = "openai_chat" + +[[components.config.profiles]] +mode = "builtin" +priority = 80 + +[components.config.profiles.builtin] +action = "redact" +detector = "email" + +[[components.config.profiles]] +mode = "local_model" +priority = 90 + +[components.config.profiles.local] +backend = "nemo_relay.pii_rampart/detector" +min_score = 0.4 +max_latency_ms = 5000 +target_path_patterns = [ + "/messages/*/content", + "/messages/*/content/*/text", + "/message", + "/message/*/text", +] +``` + +Keep contextual classifiers limited to normalized content paths. Sending model +names, tool IDs, trace IDs, routing fields, or arbitrary provider metadata to +a classifier increases false positives and exposes data outside the intended +policy boundary. + +The generic local-model payload deadline defaults to 250 ms. Contextual models +can need more time for large selected payloads. The Rampart profile above uses +5000 ms so a request near the 64 KiB provider limit has practical headroom on +typical client hardware; benchmark representative inputs on deployment +hardware before lowering it. + +The optional Rampart provider is distributed as a manifest-backed Python source +bundle under `crates/pii-redaction/providers/rampart`. Prefetch its pinned model +before adding and enabling the worker: + +```bash +cd crates/pii-redaction/providers/rampart +uvx --from . nemo-relay-pii-rampart-prefetch +nemo-relay plugins add ./relay-plugin.toml +nemo-relay plugins enable nemo_relay.pii_rampart +``` + +Activation remains offline-only. The prefetch command and activation both +verify the pinned model files before ONNX Runtime loads them. + ## Path Semantics -`target_paths` are exact JSON-pointer matches. +`target_paths` are exact JSON-pointer matches. `target_path_patterns` also +accepts `*` as one complete path segment. It does not provide recursive or +partial-segment matching. When both lists are empty, the selected backend +inspects every string leaf and reports a configuration warning. Use explicit +content paths for contextual local models so identifiers and provider metadata +do not enter the classifier accidentally. The plugin uses different payload boundaries for tools and LLMs: @@ -326,7 +438,8 @@ The plugin uses different payload boundaries for tools and LLMs: responses use the active built-in codec for normalized projection. Prefer normalized Relay paths such as: - `/headers/authorization` for an exact request header field - - `/messages/0/content` for request message content + - `/messages/*/content` for request message content + - `/messages/*/content/*/text` for multimodal request text - `/message` for the normalized assistant response text - Marks and non-tool, non-LLM scopes sanitize `data`, `category_profile`, and `metadata` as separate JSON values. A target path is evaluated independently @@ -406,23 +519,3 @@ For tool and LLM scope events, the PII redaction scope sanitizer leaves `data`, LLM sanitizer already owns that lifecycle boundary. This prevents repeated hashing or masking. Other scope categories use `input` for starts and `output` for ends. - -## Local Model Mode - -`local_model` is reserved for a future in-process local-model backend. - -### Current Status - -The current local-model status is: - -- The plugin contract accepts `mode = "local_model"`. -- The `local` section supports: - - `backend` - - `model_id` - - `detector_profile` - - `allow_network` - - `max_latency_ms` -- Actual behavior depends on a runtime-installed local backend provider. - -Without a provider, runtimes report the local backend as unavailable during -plugin initialization. diff --git a/examples/python-grpc-worker-plugin/nemo_relay_python_grpc_worker_example/worker.py b/examples/python-grpc-worker-plugin/nemo_relay_python_grpc_worker_example/worker.py index fbdfd6a2c..e9d9060bb 100644 --- a/examples/python-grpc-worker-plugin/nemo_relay_python_grpc_worker_example/worker.py +++ b/examples/python-grpc-worker-plugin/nemo_relay_python_grpc_worker_example/worker.py @@ -64,7 +64,18 @@ async def tag_tool_request(tool_name: str, args: Json) -> Json: ) return tagged_args + async def echo_local_model(request: Json) -> Json: + return { + "provider": "python_grpc_worker", + "request": request, + } + ctx.register_tool_request_intercept("tag_tool_request", tag_tool_request) + ctx.register_inference_provider( + "echo", + "examples.python_grpc_worker.echo.v1", + echo_local_model, + ) def _tag_json(value: Json, tag: str) -> Json: diff --git a/examples/python-grpc-worker-plugin/relay-plugin.toml b/examples/python-grpc-worker-plugin/relay-plugin.toml index c00d9ba80..b55df6679 100644 --- a/examples/python-grpc-worker-plugin/relay-plugin.toml +++ b/examples/python-grpc-worker-plugin/relay-plugin.toml @@ -22,7 +22,7 @@ manifest_root = "." artifact = "nemo_relay_python_grpc_worker_example/worker.py" [integrity] -sha256 = "sha256:966849be254cc6299a17a4bb65500363e9a48f98cc1e0091192e42b23821486f" +sha256 = "sha256:f7313049118e931adfd7d1c9ab18dc5aefedf35dd01d5efea4dc12306ee08358" [load] runtime = "python" diff --git a/go/nemo_relay/pii_redaction.go b/go/nemo_relay/pii_redaction.go index 871a83b60..43693abec 100644 --- a/go/nemo_relay/pii_redaction.go +++ b/go/nemo_relay/pii_redaction.go @@ -3,6 +3,8 @@ package nemo_relay +import "encoding/json" + // PiiRedactionPluginKind is the top-level plugin kind used by the built-in PII redaction component. const PiiRedactionPluginKind = "pii_redaction" @@ -18,13 +20,27 @@ type PiiRedactionBuiltinConfig struct { UnmaskedSuffix *int32 `json:"unmasked_suffix,omitempty"` } -// PiiRedactionLocalModelConfig configures the future local-model redaction backend. +// PiiRedactionLocalModelConfig configures a worker-backed local-model redaction provider. type PiiRedactionLocalModelConfig struct { - Backend string `json:"backend,omitempty"` - ModelID string `json:"model_id,omitempty"` - DetectorProfile string `json:"detector_profile,omitempty"` - AllowNetwork *bool `json:"allow_network,omitempty"` - MaxLatencyMS *int32 `json:"max_latency_ms,omitempty"` + Backend string `json:"backend,omitempty"` + ModelID string `json:"model_id,omitempty"` + DetectorProfile string `json:"detector_profile,omitempty"` + TargetPaths []string `json:"target_paths,omitempty"` + TargetPathPatterns []string `json:"target_path_patterns,omitempty"` + MinScore *float64 `json:"min_score,omitempty"` + ExcludedLabels []string `json:"excluded_labels,omitempty"` + Replacement *string `json:"replacement,omitempty"` + AllowNetwork *bool `json:"allow_network,omitempty"` + MaxLatencyMS *int32 `json:"max_latency_ms,omitempty"` +} + +// PiiRedactionProfile configures one ordered PII redaction backend. +type PiiRedactionProfile struct { + Enabled bool `json:"enabled"` + Mode string `json:"mode"` + Priority int32 `json:"priority"` + Builtin *PiiRedactionBuiltinConfig `json:"builtin,omitempty"` + Local *PiiRedactionLocalModelConfig `json:"local,omitempty"` } // PiiRedactionConfig is the canonical Go shape for the PII redaction plugin config document. @@ -38,11 +54,31 @@ type PiiRedactionConfig struct { ToolOutput bool `json:"tool_output"` Priority int32 `json:"priority,omitempty"` Codec string `json:"codec,omitempty"` + Profiles []PiiRedactionProfile `json:"profiles,omitempty"` Builtin *PiiRedactionBuiltinConfig `json:"builtin,omitempty"` Local *PiiRedactionLocalModelConfig `json:"local,omitempty"` Policy *ConfigPolicy `json:"policy,omitempty"` } +// MarshalJSON omits legacy top-level fields when profile composition is used. +func (config PiiRedactionConfig) MarshalJSON() ([]byte, error) { + if len(config.Profiles) == 0 { + type configAlias PiiRedactionConfig + return json.Marshal(configAlias(config)) + } + return json.Marshal(struct { + Version uint32 `json:"version,omitempty"` + Codec string `json:"codec,omitempty"` + Profiles []PiiRedactionProfile `json:"profiles"` + Policy *ConfigPolicy `json:"policy,omitempty"` + }{ + Version: config.Version, + Codec: config.Codec, + Profiles: config.Profiles, + Policy: config.Policy, + }) +} + // PiiRedactionComponentSpec wraps one PII redaction config as a top-level plugin component. type PiiRedactionComponentSpec struct { Enabled bool `json:"enabled,omitempty"` @@ -78,6 +114,15 @@ func NewPiiRedactionLocalModelConfig() PiiRedactionLocalModelConfig { return PiiRedactionLocalModelConfig{} } +// NewPiiRedactionProfile returns one enabled built-in profile with default priority. +func NewPiiRedactionProfile() PiiRedactionProfile { + return PiiRedactionProfile{ + Enabled: true, + Mode: "builtin", + Priority: 100, + } +} + // NewPiiRedactionComponentSpec wraps PII redaction config as an enabled component. func NewPiiRedactionComponentSpec(config PiiRedactionConfig) PiiRedactionComponentSpec { return PiiRedactionComponentSpec{ diff --git a/go/nemo_relay/pii_redaction/pii_redaction.go b/go/nemo_relay/pii_redaction/pii_redaction.go index 5a127ed96..68297995a 100644 --- a/go/nemo_relay/pii_redaction/pii_redaction.go +++ b/go/nemo_relay/pii_redaction/pii_redaction.go @@ -11,9 +11,12 @@ type Config = nemo_relay.PiiRedactionConfig // BuiltinConfig configures deterministic built-in redaction. type BuiltinConfig = nemo_relay.PiiRedactionBuiltinConfig -// LocalModelConfig configures the future local-model redaction backend. +// LocalModelConfig configures a worker-backed local-model redaction provider. type LocalModelConfig = nemo_relay.PiiRedactionLocalModelConfig +// Profile configures one ordered PII redaction backend. +type Profile = nemo_relay.PiiRedactionProfile + // ComponentSpec wraps PII redaction config as a top-level plugin component. type ComponentSpec = nemo_relay.PiiRedactionComponentSpec @@ -41,6 +44,11 @@ func NewLocalModelConfig() LocalModelConfig { return nemo_relay.NewPiiRedactionLocalModelConfig() } +// NewProfile returns one enabled built-in profile with default priority. +func NewProfile() Profile { + return nemo_relay.NewPiiRedactionProfile() +} + // NewComponentSpec wraps PII redaction config as an enabled component. func NewComponentSpec(config Config) ComponentSpec { return nemo_relay.NewPiiRedactionComponentSpec(config) diff --git a/go/nemo_relay/pii_redaction/pii_redaction_test.go b/go/nemo_relay/pii_redaction/pii_redaction_test.go index 3bb3301a2..95c62a127 100644 --- a/go/nemo_relay/pii_redaction/pii_redaction_test.go +++ b/go/nemo_relay/pii_redaction/pii_redaction_test.go @@ -27,13 +27,39 @@ func TestPiiRedactionShorthandHelpers(t *testing.T) { func TestPiiRedactionComponentSpecAndLocalModelHelpers(t *testing.T) { config := NewConfig() local := NewLocalModelConfig() - local.Backend = "local" + minScore := 0.75 + replacement := "[PRIVATE]" + allowNetwork := false + maxLatencyMS := int32(250) + local.Backend = "nemo_relay.pii_rampart/detector" local.ModelID = "pii-model" - config.Mode = "local" + local.DetectorProfile = "default" + local.TargetPaths = []string{"/message"} + local.TargetPathPatterns = []string{"/messages/*/content"} + local.MinScore = &minScore + local.ExcludedLabels = []string{"CITY"} + local.Replacement = &replacement + local.AllowNetwork = &allowNetwork + local.MaxLatencyMS = &maxLatencyMS + config.Mode = "local_model" config.Local = &local spec := NewComponentSpec(config) - if !spec.Enabled || spec.Config.Local == nil || spec.Config.Local.ModelID != "pii-model" { + if !spec.Enabled || + spec.Config.Local == nil || + spec.Config.Local.ModelID != "pii-model" || + spec.Config.Local.DetectorProfile != "default" || + len(spec.Config.Local.TargetPaths) != 1 || + len(spec.Config.Local.TargetPathPatterns) != 1 || + spec.Config.Local.MinScore == nil || + *spec.Config.Local.MinScore != minScore || + len(spec.Config.Local.ExcludedLabels) != 1 || + spec.Config.Local.Replacement == nil || + *spec.Config.Local.Replacement != replacement || + spec.Config.Local.AllowNetwork == nil || + *spec.Config.Local.AllowNetwork || + spec.Config.Local.MaxLatencyMS == nil || + *spec.Config.Local.MaxLatencyMS != maxLatencyMS { t.Fatalf("unexpected PII redaction component spec: %#v", spec) } } diff --git a/go/nemo_relay/pii_redaction_test.go b/go/nemo_relay/pii_redaction_test.go index 5b5728adc..d9591fbed 100644 --- a/go/nemo_relay/pii_redaction_test.go +++ b/go/nemo_relay/pii_redaction_test.go @@ -3,7 +3,11 @@ package nemo_relay -import "testing" +import ( + "encoding/json" + "reflect" + "testing" +) func TestPiiRedactionConfigHelpers(t *testing.T) { config := NewPiiRedactionConfig() @@ -18,9 +22,13 @@ func TestPiiRedactionConfigHelpers(t *testing.T) { t.Fatalf("unexpected built-in redaction defaults: %#v", builtin) } local := NewPiiRedactionLocalModelConfig() - if local != (PiiRedactionLocalModelConfig{}) { + if !reflect.DeepEqual(local, PiiRedactionLocalModelConfig{}) { t.Fatalf("unexpected local model defaults: %#v", local) } + profile := NewPiiRedactionProfile() + if !profile.Enabled || profile.Mode != "builtin" || profile.Priority != 100 { + t.Fatalf("unexpected profile defaults: %#v", profile) + } config.Builtin = &builtin component := PiiRedactionComponent(config) @@ -42,6 +50,30 @@ func TestPiiRedactionConfigHelpers(t *testing.T) { } } +func TestPiiRedactionProfilesOmitLegacyTopLevelFields(t *testing.T) { + config := NewPiiRedactionConfig() + config.Profiles = []PiiRedactionProfile{ + NewPiiRedactionProfile(), + } + serialized, err := json.Marshal(config) + if err != nil { + t.Fatalf("marshal profile config: %v", err) + } + var value map[string]any + if err := json.Unmarshal(serialized, &value); err != nil { + t.Fatalf("decode profile config: %v", err) + } + if _, present := value["mode"]; present { + t.Fatalf("profile config retained legacy mode: %#v", value) + } + if _, present := value["input"]; present { + t.Fatalf("profile config retained legacy input: %#v", value) + } + if len(value["profiles"].([]any)) != 1 { + t.Fatalf("unexpected profile config: %#v", value) + } +} + func TestPiiRedactionValidationRejectsBadValues(t *testing.T) { config := NewPiiRedactionConfig() config.Input = false diff --git a/justfile b/justfile index c64a671ab..937dec95d 100644 --- a/justfile +++ b/justfile @@ -971,6 +971,9 @@ check-python-worker-proto: } assert pb.SUBSCRIBER == 1 assert pb.LLM_STREAM_EXECUTION_INTERCEPT == 25 + assert pb.INFERENCE_PROVIDER == 40 + assert pb.Registration.DESCRIPTOR.fields_by_name["contract"].number == 5 + assert pb.InvokeRequest.DESCRIPTOR.fields_by_name["provider"].number == 13 PY generate-worker-plugin-lockfile: diff --git a/python/nemo_relay/pii_redaction.py b/python/nemo_relay/pii_redaction.py index 53c2b77f8..de4f443d0 100644 --- a/python/nemo_relay/pii_redaction.py +++ b/python/nemo_relay/pii_redaction.py @@ -103,11 +103,16 @@ def to_dict(self) -> JsonObject: @dataclass(slots=True) class LocalModelConfig: - """Future local-model backend seam settings.""" + """Worker-backed local-model redaction settings.""" backend: str | None = None model_id: str | None = None detector_profile: str | None = None + target_paths: list[str] = field(default_factory=list) + target_path_patterns: list[str] = field(default_factory=list) + min_score: float | None = None + excluded_labels: list[str] = field(default_factory=list) + replacement: str | None = None allow_network: bool | None = None max_latency_ms: int | None = None @@ -118,12 +123,40 @@ def to_dict(self) -> JsonObject: "backend": self.backend, "model_id": self.model_id, "detector_profile": self.detector_profile, + "target_paths": self.target_paths or None, + "target_path_patterns": self.target_path_patterns or None, + "min_score": self.min_score, + "excluded_labels": self.excluded_labels or None, + "replacement": self.replacement, "allow_network": self.allow_network, "max_latency_ms": self.max_latency_ms, } ) +@dataclass(slots=True) +class PiiRedactionProfile: + """One ordered PII redaction backend profile.""" + + enabled: bool = True + mode: Literal["builtin", "local_model"] = "builtin" + priority: int = 100 + builtin: BuiltinConfig | None = None + local: LocalModelConfig | None = None + + def to_dict(self) -> JsonObject: + """Serialize this profile to the canonical JSON object shape.""" + return _normalize_object( + { + "enabled": self.enabled, + "mode": self.mode, + "priority": self.priority, + "builtin": self.builtin, + "local": self.local, + } + ) + + @dataclass(slots=True) class PiiRedactionConfig: """Canonical config document for the top-level PII redaction component.""" @@ -137,12 +170,22 @@ class PiiRedactionConfig: mark: bool = True priority: int = 100 codec: Literal["openai_chat", "openai_responses", "anthropic_messages"] | str | None = None + profiles: list[PiiRedactionProfile] = field(default_factory=list) builtin: BuiltinConfig | None = None local: LocalModelConfig | None = None policy: ConfigPolicy = field(default_factory=ConfigPolicy) def to_dict(self) -> JsonObject: """Serialize this PII redaction config to the canonical JSON object shape.""" + if self.profiles: + return _normalize_object( + { + "version": self.version, + "codec": self.codec, + "profiles": self.profiles, + "policy": self.policy, + } + ) return _normalize_object( { "version": self.version, @@ -187,7 +230,7 @@ def validate_config(config: PiiRedactionConfig | JsonObject) -> ConfigReport: components=[ComponentSpec(config)], ) ) - return cast(ConfigReport, report) + return report __all__ = [ @@ -199,5 +242,6 @@ def validate_config(config: PiiRedactionConfig | JsonObject) -> ConfigReport: "LocalModelConfig", "PII_REDACTION_PLUGIN_KIND", "PiiRedactionConfig", + "PiiRedactionProfile", "validate_config", ] diff --git a/python/nemo_relay/pii_redaction.pyi b/python/nemo_relay/pii_redaction.pyi index 244f6a3ef..fd6abaa72 100644 --- a/python/nemo_relay/pii_redaction.pyi +++ b/python/nemo_relay/pii_redaction.pyi @@ -44,10 +44,24 @@ class LocalModelConfig: backend: str | None = ... model_id: str | None = ... detector_profile: str | None = ... + target_paths: list[str] = field(default_factory=list) + target_path_patterns: list[str] = field(default_factory=list) + min_score: float | None = ... + excluded_labels: list[str] = field(default_factory=list) + replacement: str | None = ... allow_network: bool | None = ... max_latency_ms: int | None = ... def to_dict(self) -> JsonObject: ... +@dataclass(slots=True) +class PiiRedactionProfile: + enabled: bool = ... + mode: Literal["builtin", "local_model"] = ... + priority: int = ... + builtin: BuiltinConfig | None = ... + local: LocalModelConfig | None = ... + def to_dict(self) -> JsonObject: ... + @dataclass(slots=True) class PiiRedactionConfig: version: int = ... @@ -59,6 +73,7 @@ class PiiRedactionConfig: mark: bool = ... priority: int = ... codec: Literal["openai_chat", "openai_responses", "anthropic_messages"] | str | None = ... + profiles: list[PiiRedactionProfile] = field(default_factory=list) builtin: BuiltinConfig | None = ... local: LocalModelConfig | None = ... policy: ConfigPolicy = field(default_factory=ConfigPolicy) diff --git a/python/plugin/README.md b/python/plugin/README.md index 051383f8d..3eac28c7a 100644 --- a/python/plugin/README.md +++ b/python/plugin/README.md @@ -26,7 +26,8 @@ protocol. - **Isolate plugin dependencies**: Run custom policy, middleware, or exporter code outside the Relay host process. - **Use the shared runtime contract**: Register subscribers, guardrails, and - intercepts through `WorkerPlugin` and `PluginContext`. + intercepts or inference providers through `WorkerPlugin` and + `PluginContext`. - **Call back into Relay safely**: Emit marks, create scopes, and continue managed execution through the host runtime handle. - **Keep worker lifecycle managed**: Let Relay provision the worker environment, @@ -104,6 +105,31 @@ worker process. For a complete manifest and runnable plugin, see the [Python gRPC worker plugin example](https://github.com/NVIDIA/NeMo-Relay/blob/main/examples/python-grpc-worker-plugin/README.md). +## Inference Providers + +Use `register_inference_provider` when a first-party Relay component owns a +versioned request-response contract and needs isolated model inference: + +```python +async def detect(request: Json) -> Json: + return { + "version": 1, + "detections": await model.detect(request["texts"]), + } + + +ctx.register_inference_provider( + "detector", + "acme.pii_detection.v1", + detect, +) +``` + +Relay publishes this provider as `/detector`. The callback may be +synchronous or asynchronous and should perform inference only. The consuming +host component selects the exact contract and owns the payload schema, +deadline, field traversal, output validation, and result application. + ## Request Intercepts LLM request intercepts return one canonical outcome: diff --git a/python/plugin/src/nemo_relay_plugin/__init__.py b/python/plugin/src/nemo_relay_plugin/__init__.py index 54466a5b5..a3bb23893 100644 --- a/python/plugin/src/nemo_relay_plugin/__init__.py +++ b/python/plugin/src/nemo_relay_plugin/__init__.py @@ -36,6 +36,7 @@ LlmOptimizationTokens: Explicit token evidence by category. LlmOptimizationTokenImpact: Baseline, effective, and saved token evidence. LlmRequestInterceptOutcome: Canonical LLM request-intercept result. + InferenceProviderCallback: Versioned inference-provider callback. ToolExecutionInterceptOutcome: Canonical tool execution-intercept result. DiagnosticLevel: Severity of a configuration diagnostic. ConfigDiagnostic: Structured configuration warning or error. @@ -75,6 +76,7 @@ Event, EventSanitizeCallback, EventSanitizeFields, + InferenceProviderCallback, Json, LlmCodecIdentity, LlmConditionalCallback, @@ -142,6 +144,7 @@ "LlmSanitizeResponseCallback", "LlmStreamNext", "LlmStreamExecutionCallback", + "InferenceProviderCallback", "PluginContext", "PluginRuntime", "PendingMarkSpec", diff --git a/python/plugin/src/nemo_relay_plugin/_api.py b/python/plugin/src/nemo_relay_plugin/_api.py index 7f8b6ccb8..d14d390ff 100644 --- a/python/plugin/src/nemo_relay_plugin/_api.py +++ b/python/plugin/src/nemo_relay_plugin/_api.py @@ -845,6 +845,7 @@ def register(self, ctx: PluginContext, config: Json) -> None | Awaitable[None]: [str, LlmRequest, "LlmStreamNext"], Iterable[Json] | AsyncIterator[Json] | Awaitable[Iterable[Json] | AsyncIterator[Json]], ] +InferenceProviderCallback: TypeAlias = Callable[[Json], Json | Awaitable[Json]] @dataclass(slots=True) @@ -865,6 +866,7 @@ class _Handlers: llm_requests: dict[str, LlmRequestCallback] llm_executions: dict[str, LlmExecutionCallback] llm_stream_executions: dict[str, LlmStreamExecutionCallback] + inference_providers: dict[str, InferenceProviderCallback] @classmethod def empty(cls) -> _Handlers: @@ -885,6 +887,7 @@ def empty(cls) -> _Handlers: llm_requests={}, llm_executions={}, llm_stream_executions={}, + inference_providers={}, ) @@ -953,6 +956,28 @@ def register_subscriber(self, name: str, callback: SubscriberCallback) -> None: self._push_registration(name, pb.SUBSCRIBER, 0, False) self._handlers.subscribers[name] = callback + def register_inference_provider( + self, + name: str, + contract: str, + callback: InferenceProviderCallback, + ) -> None: + """Register a named inference provider for a versioned host contract. + + Args: + name: Stable provider name selected by a consuming host component. + contract: Versioned request-response contract implemented by the provider. + callback: Function receiving and returning component-owned JSON. + The callback can return a value directly or through an + awaitable. + + Provider boundary: + Providers perform model inference only. The consuming host + component owns field selection, policy, and output application. + """ + self._push_registration(name, pb.INFERENCE_PROVIDER, 0, False, contract=contract) + self._handlers.inference_providers[name] = callback + def _register_event_sanitizer( self, name: str, @@ -1260,7 +1285,15 @@ def register_llm_stream_execution_intercept( self._push_registration(name, pb.LLM_STREAM_EXECUTION_INTERCEPT, priority, False) self._handlers.llm_stream_executions[name] = callback - def _push_registration(self, name: str, surface: int, priority: int, break_chain: bool) -> None: + def _push_registration( + self, + name: str, + surface: int, + priority: int, + break_chain: bool, + *, + contract: str = "", + ) -> None: if any( registration.local_name == name and registration.surface == surface for registration in self._handlers.registrations @@ -1272,6 +1305,7 @@ def _push_registration(self, name: str, surface: int, priority: int, break_chain surface=surface, priority=priority, break_chain=break_chain, + contract=contract, ) ) @@ -2060,6 +2094,19 @@ async def _invoke_result(self, request: Any) -> Any: ), ) ) + if request.surface == pb.INFERENCE_PROVIDER: + result = await _maybe_await( + self._handler( + self._handlers.inference_providers, + request.registration_name, + )( + _decode_required_envelope( + request.provider, + "inference provider request", + ) + ) + ) + return _json_response(result) raise WorkerSdkError(f"unsupported registration surface {request.surface}") async def _invoke_llm_result(self, request: Any) -> Any: @@ -2195,6 +2242,7 @@ def _all_surfaces() -> list[int]: pb.LLM_REQUEST_INTERCEPT, pb.LLM_EXECUTION_INTERCEPT, pb.LLM_STREAM_EXECUTION_INTERCEPT, + pb.INFERENCE_PROVIDER, ] diff --git a/python/tests/plugin/test_public_api_docstrings.py b/python/tests/plugin/test_public_api_docstrings.py index 31abf7427..4ab3074c8 100644 --- a/python/tests/plugin/test_public_api_docstrings.py +++ b/python/tests/plugin/test_public_api_docstrings.py @@ -35,6 +35,7 @@ "LlmRequestCallback", "LlmExecutionCallback", "LlmStreamExecutionCallback", + "InferenceProviderCallback", } diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py index 7ab918d20..20062585d 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -363,6 +363,9 @@ async def llm_stream_execution(name: str, request: Json, next_call: Any) -> Asyn async for chunk in stream: yield _tag(chunk, "llm_stream_execution") + async def inference_provider(request: Json) -> Json: + return _tag(request, "local_model") + ctx.register_subscriber("subscriber", subscriber) ctx.register_mark_sanitize_guardrail("event_sanitize", mark_sanitize, priority=1) ctx.register_scope_sanitize_start_guardrail("event_sanitize", scope_start_sanitize, priority=2) @@ -378,6 +381,11 @@ async def llm_stream_execution(name: str, request: Json, next_call: Any) -> Asyn ctx.register_llm_request_intercept("llm_request", llm_request, priority=9, break_chain=True) ctx.register_llm_execution_intercept("llm_execution", llm_execution, priority=10) ctx.register_llm_stream_execution_intercept("llm_stream_execution", llm_stream_execution, priority=11) + ctx.register_inference_provider( + "local_model", + "test.echo.v1", + inference_provider, + ) @pytest.fixture(name="host_stub") @@ -411,6 +419,9 @@ def test_generated_proto_matches_worker_contract(): assert pb.MARK_SANITIZE_GUARDRAIL == 30 assert pb.SCOPE_SANITIZE_START_GUARDRAIL == 31 assert pb.SCOPE_SANITIZE_END_GUARDRAIL == 32 + assert pb.INFERENCE_PROVIDER == 40 + assert pb.Registration.DESCRIPTOR.fields_by_name["contract"].number == 5 + assert pb.InvokeRequest.DESCRIPTOR.fields_by_name["provider"].number == 13 assert pb.CUSTOM == 10 @@ -444,25 +455,32 @@ async def test_health_handshake_validate_register_and_all_surfaces(service: _Wor register = await _register(service) registrations = [ - (registration.local_name, registration.surface, registration.priority, registration.break_chain) + ( + registration.local_name, + registration.surface, + registration.priority, + registration.break_chain, + registration.contract, + ) for registration in register.registrations ] assert registrations == [ - ("subscriber", pb.SUBSCRIBER, 0, False), - ("event_sanitize", pb.MARK_SANITIZE_GUARDRAIL, 1, False), - ("event_sanitize", pb.SCOPE_SANITIZE_START_GUARDRAIL, 2, False), - ("scope_end_sanitize", pb.SCOPE_SANITIZE_END_GUARDRAIL, 3, False), - ("tool_sanitize", pb.TOOL_SANITIZE_REQUEST_GUARDRAIL, 1, False), - ("tool_sanitize", pb.TOOL_SANITIZE_RESPONSE_GUARDRAIL, 2, False), - ("tool_conditional", pb.TOOL_CONDITIONAL_EXECUTION_GUARDRAIL, 3, False), - ("tool_request", pb.TOOL_REQUEST_INTERCEPT, 4, True), - ("tool_execution", pb.TOOL_EXECUTION_INTERCEPT, 5, False), - ("llm_sanitize_request", pb.LLM_SANITIZE_REQUEST_GUARDRAIL, 6, False), - ("llm_sanitize_response", pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, 7, False), - ("llm_conditional", pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, 8, False), - ("llm_request", pb.LLM_REQUEST_INTERCEPT, 9, True), - ("llm_execution", pb.LLM_EXECUTION_INTERCEPT, 10, False), - ("llm_stream_execution", pb.LLM_STREAM_EXECUTION_INTERCEPT, 11, False), + ("subscriber", pb.SUBSCRIBER, 0, False, ""), + ("event_sanitize", pb.MARK_SANITIZE_GUARDRAIL, 1, False, ""), + ("event_sanitize", pb.SCOPE_SANITIZE_START_GUARDRAIL, 2, False, ""), + ("scope_end_sanitize", pb.SCOPE_SANITIZE_END_GUARDRAIL, 3, False, ""), + ("tool_sanitize", pb.TOOL_SANITIZE_REQUEST_GUARDRAIL, 1, False, ""), + ("tool_sanitize", pb.TOOL_SANITIZE_RESPONSE_GUARDRAIL, 2, False, ""), + ("tool_conditional", pb.TOOL_CONDITIONAL_EXECUTION_GUARDRAIL, 3, False, ""), + ("tool_request", pb.TOOL_REQUEST_INTERCEPT, 4, True, ""), + ("tool_execution", pb.TOOL_EXECUTION_INTERCEPT, 5, False, ""), + ("llm_sanitize_request", pb.LLM_SANITIZE_REQUEST_GUARDRAIL, 6, False, ""), + ("llm_sanitize_response", pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, 7, False, ""), + ("llm_conditional", pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, 8, False, ""), + ("llm_request", pb.LLM_REQUEST_INTERCEPT, 9, True, ""), + ("llm_execution", pb.LLM_EXECUTION_INTERCEPT, 10, False, ""), + ("llm_stream_execution", pb.LLM_STREAM_EXECUTION_INTERCEPT, 11, False, ""), + ("local_model", pb.INFERENCE_PROVIDER, 0, False, "test.echo.v1"), ] @@ -1438,6 +1456,16 @@ async def test_unary_invoke_success_paths(service: _WorkerService, host_stub: Re assert llm_execution["tag"] == "llm_execution" assert llm_execution["next_llm"]["content"]["llm_execute_gpt-test"] + local_model = await service.Invoke( + _provider_request("local_model", {"text": "private"}), + AbortContext(), + ) + assert local_model.WhichOneof("result") == "json" + assert _envelope_value(local_model.json.value) == { + "text": "private", + "tag": "local_model", + } + async def test_unary_invoke_failure_paths(service: _WorkerService): await _register(service) @@ -1452,6 +1480,13 @@ async def test_unary_invoke_failure_paths(service: _WorkerService): assert missing_handler.WhichOneof("result") == "error" assert "not registered" in missing_handler.error.message + missing_provider = await service.Invoke( + _provider_request("missing", {}), + AbortContext(), + ) + assert missing_provider.WhichOneof("result") == "error" + assert "not registered" in missing_provider.error.message + unsupported = await service.Invoke( _tool_request("tool_request", pb.REGISTRATION_SURFACE_UNSPECIFIED, {}), AbortContext(), @@ -2738,6 +2773,15 @@ def _tool_request(registration_name: str, surface: int, value: Json) -> Any: ) +def _provider_request(registration_name: str, value: Json) -> Any: + return _invoke_request( + registration_name, + pb.INFERENCE_PROVIDER, + continuation_id="", + provider=_json_envelope(JSON_SCHEMA, value), + ) + + def _llm_payload( *, model_name: str = "model", @@ -2826,4 +2870,5 @@ def _all_expected_surfaces() -> list[int]: pb.LLM_REQUEST_INTERCEPT, pb.LLM_EXECUTION_INTERCEPT, pb.LLM_STREAM_EXECUTION_INTERCEPT, + pb.INFERENCE_PROVIDER, ] diff --git a/python/tests/test_pii_redaction_plugin.py b/python/tests/test_pii_redaction_plugin.py index cbed8c4c2..ab57329b3 100644 --- a/python/tests/test_pii_redaction_plugin.py +++ b/python/tests/test_pii_redaction_plugin.py @@ -13,6 +13,7 @@ ConfigPolicy, LocalModelConfig, PiiRedactionConfig, + PiiRedactionProfile, validate_config, ) @@ -29,6 +30,29 @@ def test_defaults_and_component_wrapper(self): "unsupported_value": "error", } assert LocalModelConfig().to_dict() == {} + assert LocalModelConfig( + backend="acme.pii/detector", + model_id="pii-model-v1", + detector_profile="default", + target_paths=["/message"], + target_path_patterns=["/messages/*/content"], + min_score=0.6, + excluded_labels=["CITY"], + replacement="[PRIVATE]", + allow_network=False, + max_latency_ms=250, + ).to_dict() == { + "backend": "acme.pii/detector", + "model_id": "pii-model-v1", + "detector_profile": "default", + "target_paths": ["/message"], + "target_path_patterns": ["/messages/*/content"], + "min_score": 0.6, + "excluded_labels": ["CITY"], + "replacement": "[PRIVATE]", + "allow_network": False, + "max_latency_ms": 250, + } wrapped = ComponentSpec(PiiRedactionConfig()).to_dict() assert wrapped["kind"] == PII_REDACTION_PLUGIN_KIND @@ -42,6 +66,36 @@ def test_defaults_and_component_wrapper(self): opted_out = PiiRedactionConfig(mark=False).to_dict() assert opted_out["mark"] is False + def test_profile_config_omits_legacy_top_level_fields(self): + config = PiiRedactionConfig( + codec="openai_chat", + profiles=[ + PiiRedactionProfile( + mode="builtin", + builtin=BuiltinConfig(detector="email"), + ), + PiiRedactionProfile( + mode="local_model", + priority=110, + local=LocalModelConfig( + backend="acme.pii/detector", + target_path_patterns=["/messages/*/content"], + ), + ), + ], + ).to_dict() + + profiles = config["profiles"] + assert isinstance(profiles, list) + local_profile = profiles[1] + assert isinstance(local_profile, dict) + local = local_profile["local"] + assert isinstance(local, dict) + assert local["backend"] == "acme.pii/detector" + assert "mode" not in config + assert "input" not in config + assert validate_config(config)["diagnostics"] == [] + def test_validation_rejects_bad_values(self): report = validate_config( PiiRedactionConfig(