diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index 6d5c0dba83..72c7224871 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -28,7 +28,7 @@ __test-cli-port-override = [] __skip-implicit-config = ["nemo-relay/__skip-implicit-config"] [dependencies] -nemo-relay = { workspace = true, features = ["guardrails-remote", "object-store", "worker-grpc"] } +nemo-relay = { workspace = true, features = ["object-store", "worker-grpc"] } nemo-relay-adaptive = { workspace = true, features = ["redis-backend"] } nemo-relay-pii-redaction.workspace = true async-stream = "0.3" diff --git a/crates/cli/src/plugins/editor_model.rs b/crates/cli/src/plugins/editor_model.rs index f7b6add268..f1c57f9177 100644 --- a/crates/cli/src/plugins/editor_model.rs +++ b/crates/cli/src/plugins/editor_model.rs @@ -19,20 +19,6 @@ use serde_json::{Map, Value, json}; use crate::error::CliError; -#[allow( - deprecated, - reason = "the CLI must edit existing Guardrails configuration until the built-in plugin is removed" -)] -mod guardrails_compat { - pub(super) type Config = nemo_relay::plugins::nemo_guardrails::component::NeMoGuardrailsConfig; - pub(super) const PLUGIN_KIND: &str = - nemo_relay::plugins::nemo_guardrails::component::NEMO_GUARDRAILS_PLUGIN_KIND; -} - -use guardrails_compat::{ - Config as NeMoGuardrailsConfig, PLUGIN_KIND as NEMO_GUARDRAILS_PLUGIN_KIND, -}; - pub(super) const POLICY_SECTION: &str = "policy"; #[derive(Debug, Clone)] @@ -50,7 +36,6 @@ pub(super) enum EditableComponent { Observability(Box>), ResourceMetrics(Box>), Adaptive(Box>), - NemoGuardrails(Box>), PiiRedaction(Box>), } @@ -60,7 +45,6 @@ impl EditableComponent { Self::Observability(_) => "Observability", Self::ResourceMetrics(_) => "Resource Metrics", Self::Adaptive(_) => "Adaptive", - Self::NemoGuardrails(_) => "NeMo Guardrails (Deprecated)", Self::PiiRedaction(_) => "PII Redaction", } } @@ -70,7 +54,6 @@ impl EditableComponent { Self::Observability(_) => ObservabilityConfig::editor_schema().fields, Self::ResourceMetrics(_) => ResourceMetricsConfig::editor_schema().fields, Self::Adaptive(_) => AdaptiveConfig::editor_schema().fields, - Self::NemoGuardrails(_) => NeMoGuardrailsConfig::editor_schema().fields, Self::PiiRedaction(_) => PiiRedactionConfig::editor_schema().fields, } } @@ -80,7 +63,6 @@ impl EditableComponent { Self::Observability(state) => state.enabled, Self::ResourceMetrics(state) => state.enabled, Self::Adaptive(state) => state.enabled, - Self::NemoGuardrails(state) => state.enabled, Self::PiiRedaction(state) => state.enabled, } } @@ -90,7 +72,6 @@ impl EditableComponent { Self::Observability(state) => state.toggle_enabled(), Self::ResourceMetrics(state) => state.toggle_enabled(), Self::Adaptive(state) => state.toggle_enabled(), - Self::NemoGuardrails(state) => state.toggle_enabled(), Self::PiiRedaction(state) => state.toggle_enabled(), } } @@ -100,7 +81,6 @@ impl EditableComponent { Self::Observability(state) => state.set_enabled(enabled), Self::ResourceMetrics(state) => state.set_enabled(enabled), Self::Adaptive(state) => state.set_enabled(enabled), - Self::NemoGuardrails(state) => state.set_enabled(enabled), Self::PiiRedaction(state) => state.set_enabled(enabled), } } @@ -110,7 +90,6 @@ impl EditableComponent { Self::Observability(state) => state.reset_enabled(), Self::ResourceMetrics(state) => state.reset_enabled(), Self::Adaptive(state) => state.reset_enabled(), - Self::NemoGuardrails(state) => state.reset_enabled(), Self::PiiRedaction(state) => state.reset_enabled(), } } @@ -120,7 +99,6 @@ impl EditableComponent { Self::Observability(state) => observability_summary(state), Self::ResourceMetrics(state) => resource_metrics_summary(state), Self::Adaptive(state) => adaptive_summary(state), - Self::NemoGuardrails(state) => nemo_guardrails_summary(state), Self::PiiRedaction(state) => pii_redaction_summary(state), } } @@ -132,9 +110,6 @@ impl EditableComponent { config_field_configured(&state.config, field).unwrap_or(false) } Self::Adaptive(state) => config_field_configured(&state.config, field).unwrap_or(false), - Self::NemoGuardrails(state) => { - config_field_configured(&state.config, field).unwrap_or(false) - } Self::PiiRedaction(state) => { config_field_configured(&state.config, field).unwrap_or(false) } @@ -155,10 +130,6 @@ impl EditableComponent { reset_config_field(&mut state.config, field)?; state.mark_config_touched(); } - Self::NemoGuardrails(state) => { - reset_config_field(&mut state.config, field)?; - state.mark_config_touched(); - } Self::PiiRedaction(state) => { reset_config_field(&mut state.config, field)?; state.mark_config_touched(); @@ -184,10 +155,6 @@ impl EditableComponent { remove_struct_field(&mut state.config, field.name)?; state.mark_config_touched(); } - Self::NemoGuardrails(state) => { - remove_struct_field(&mut state.config, field.name)?; - state.mark_config_touched(); - } Self::PiiRedaction(state) => { remove_struct_field(&mut state.config, field.name)?; state.mark_config_touched(); @@ -201,7 +168,6 @@ impl EditableComponent { Self::Observability(state) => store_observability_state(config, state), Self::ResourceMetrics(state) => store_resource_metrics_state(config, state), Self::Adaptive(state) => store_adaptive_state(config, state), - Self::NemoGuardrails(state) => store_nemo_guardrails_state(config, state), Self::PiiRedaction(state) => store_pii_redaction_state(config, state), } } @@ -230,7 +196,6 @@ pub(super) fn editable_components( EditableComponent::Observability(Box::new(component_observability_state(config)?)), EditableComponent::ResourceMetrics(Box::new(component_resource_metrics_state(config)?)), EditableComponent::Adaptive(Box::new(component_adaptive_state(config)?)), - EditableComponent::NemoGuardrails(Box::new(component_nemo_guardrails_state(config)?)), EditableComponent::PiiRedaction(Box::new(component_pii_redaction_state(config)?)), ]; Ok(components) @@ -432,12 +397,6 @@ pub(super) fn component_adaptive_state( component_editor_state(config, ADAPTIVE_PLUGIN_KIND, false) } -pub(super) fn component_nemo_guardrails_state( - config: &PluginConfig, -) -> Result, CliError> { - component_editor_state(config, NEMO_GUARDRAILS_PLUGIN_KIND, false) -} - pub(super) fn component_pii_redaction_state( config: &PluginConfig, ) -> Result, CliError> { @@ -492,22 +451,6 @@ pub(super) fn store_adaptive_state( Ok(()) } -pub(super) fn store_nemo_guardrails_state( - config: &mut PluginConfig, - state: &ComponentEditorState, -) -> Result<(), CliError> { - if state.should_store(state.config_touched || nemo_guardrails_configured(&state.config)) { - store_component_editor_config( - config, - NEMO_GUARDRAILS_PLUGIN_KIND, - state.enabled, - nemo_guardrails_config_map(&state.config)?, - merge_nemo_guardrails_editor_config, - ); - } - Ok(()) -} - pub(super) fn store_pii_redaction_state( config: &mut PluginConfig, state: &ComponentEditorState, @@ -850,23 +793,6 @@ pub(super) fn adaptive_config_map(config: &AdaptiveConfig) -> Result Result, CliError> { - let value = serde_json::to_value(config).map_err(serde_error)?; - match value { - Value::Object(mut map) => { - if is_version_one(map.get("version")) { - map.remove("version"); - } - Ok(map) - } - _ => Err(CliError::Config( - "nemo_guardrails config must serialize to an object".into(), - )), - } -} - pub(super) fn pii_redaction_config_map( config: &PiiRedactionConfig, ) -> Result, CliError> { @@ -919,21 +845,6 @@ pub(super) fn merge_adaptive_editor_config( ); } -pub(super) fn merge_nemo_guardrails_editor_config( - existing: &mut Map, - edited: Map, -) { - if is_version_one(existing.get("version")) { - existing.remove("version"); - } - merge_known_editor_object( - existing, - edited, - &nested_editor_keys(NeMoGuardrailsConfig::editor_schema()), - NeMoGuardrailsConfig::editor_schema(), - ); -} - pub(super) fn merge_pii_redaction_editor_config( existing: &mut Map, edited: Map, @@ -1171,35 +1082,6 @@ pub(super) fn adaptive_summary(state: &ComponentEditorState) -> ) } -pub(super) fn nemo_guardrails_configured(config: &NeMoGuardrailsConfig) -> bool { - NeMoGuardrailsConfig::editor_schema() - .fields - .iter() - .filter(|field| field.name != POLICY_SECTION) - .any(|field| config_field_configured(config, *field).unwrap_or(false)) -} - -pub(super) fn nemo_guardrails_summary( - state: &ComponentEditorState, -) -> String { - let configured_fields = NeMoGuardrailsConfig::editor_schema() - .fields - .iter() - .filter(|field| field.name != POLICY_SECTION) - .filter(|field| config_field_configured(&state.config, **field).unwrap_or(false)) - .map(|field| field.label) - .collect::>(); - format!( - "component {}, fields {}", - if state.enabled { "enabled" } else { "disabled" }, - if configured_fields.is_empty() { - "none".into() - } else { - configured_fields.join(", ") - } - ) -} - pub(super) fn pii_redaction_configured(config: &PiiRedactionConfig) -> bool { PiiRedactionConfig::editor_schema() .fields diff --git a/crates/cli/src/plugins/prompt.rs b/crates/cli/src/plugins/prompt.rs index b510fefaf2..b1c31eed8e 100644 --- a/crates/cli/src/plugins/prompt.rs +++ b/crates/cli/src/plugins/prompt.rs @@ -175,10 +175,6 @@ fn edit_component_field( edit_config_field(theme, &mut state.config, field)?; state.mark_config_touched(); } - EditableComponent::NemoGuardrails(state) => { - edit_config_field(theme, &mut state.config, field)?; - state.mark_config_touched(); - } EditableComponent::PiiRedaction(state) => { edit_config_field(theme, &mut state.config, field)?; state.mark_config_touched(); diff --git a/crates/cli/src/server/mod.rs b/crates/cli/src/server/mod.rs index 998f109efd..0c8279a3d6 100644 --- a/crates/cli/src/server/mod.rs +++ b/crates/cli/src/server/mod.rs @@ -999,6 +999,7 @@ pub(crate) struct ServerPluginActivation { } const REMOVED_SWITCHYARD_MESSAGE: &str = "the built-in Switchyard service integration was removed in NeMo Relay >=0.8.0; remove this `[[components]]` entry and refer to the NeMo Relay migration guides for current Switchyard migration information: https://docs.nvidia.com/nemo/relay/reference/migration-guides"; +const REMOVED_NEMO_GUARDRAILS_MESSAGE: &str = "the built-in NeMo Guardrails integration was removed in NeMo Relay >=0.10.0; remove this `[[components]]` entry and refer to the migration guide: https://docs.nvidia.com/nemo/relay/reference/migration-guides#remove-the-built-in-nemo-guardrails-component"; impl ServerPluginActivation { pub(crate) fn clear(mut self) -> Result<(), CliError> { @@ -1012,6 +1013,7 @@ impl ServerPluginActivation { pub(crate) enum PluginComponentSetupError { Adaptive(String), PiiRedaction(String), + RemovedNemoGuardrails, RemovedSwitchyard, } @@ -1020,6 +1022,7 @@ impl PluginComponentSetupError { match self { Self::Adaptive(_) => "Adaptive plugin", Self::PiiRedaction(_) => "PII redaction plugin", + Self::RemovedNemoGuardrails => "NeMo Guardrails migration", Self::RemovedSwitchyard => "Switchyard migration", } } @@ -1029,6 +1032,7 @@ impl PluginComponentSetupError { Self::Adaptive(error) | Self::PiiRedaction(error) => { format!("registration failed: {error}") } + Self::RemovedNemoGuardrails => REMOVED_NEMO_GUARDRAILS_MESSAGE.into(), Self::RemovedSwitchyard => REMOVED_SWITCHYARD_MESSAGE.into(), } } @@ -1046,6 +1050,7 @@ impl std::fmt::Display for PluginComponentSetupError { "PII redaction plugin registration failed: {error}" ) } + Self::RemovedNemoGuardrails => formatter.write_str(REMOVED_NEMO_GUARDRAILS_MESSAGE), Self::RemovedSwitchyard => formatter.write_str(REMOVED_SWITCHYARD_MESSAGE), } } @@ -1068,6 +1073,13 @@ pub(crate) fn register_and_validate_plugin_components( { errors.push(PluginComponentSetupError::RemovedSwitchyard); } + if plugin_config + .components + .iter() + .any(|component| component.kind == "nemo_guardrails") + { + errors.push(PluginComponentSetupError::RemovedNemoGuardrails); + } errors } diff --git a/crates/cli/tests/coverage/shared/plugins_tests.rs b/crates/cli/tests/coverage/shared/plugins_tests.rs index 8f90c633d9..7e5d44f9c7 100644 --- a/crates/cli/tests/coverage/shared/plugins_tests.rs +++ b/crates/cli/tests/coverage/shared/plugins_tests.rs @@ -18,25 +18,6 @@ use nemo_relay_pii_redaction::component::{PII_REDACTION_PLUGIN_KIND, PiiRedactio use serde_json::Map; use std::path::PathBuf; -#[allow( - deprecated, - reason = "compatibility tests cover the built-in Guardrails editor until its scheduled removal" -)] -mod guardrails_compat { - pub(super) type Config = nemo_relay::plugins::nemo_guardrails::component::NeMoGuardrailsConfig; - pub(super) type LocalConfig = - nemo_relay::plugins::nemo_guardrails::component::LocalBackendConfig; - pub(super) type RemoteConfig = - nemo_relay::plugins::nemo_guardrails::component::RemoteBackendConfig; - pub(super) const PLUGIN_KIND: &str = - nemo_relay::plugins::nemo_guardrails::component::NEMO_GUARDRAILS_PLUGIN_KIND; -} - -use guardrails_compat::{ - Config as NeMoGuardrailsConfig, LocalConfig as LocalBackendConfig, - PLUGIN_KIND as NEMO_GUARDRAILS_PLUGIN_KIND, RemoteConfig as RemoteBackendConfig, -}; - #[test] fn plugin_document_policy_edits_preserve_host_extensions_and_prune_default_fields() { use nemo_relay::plugin::UnsupportedBehavior; @@ -265,54 +246,6 @@ fn adaptive_component_config(agent_id: &str) -> serde_json::Map { .clone() } -fn guardrails_component_config(config_id: &str) -> serde_json::Map { - json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": config_id - } - }) - .as_object() - .unwrap() - .clone() -} - -fn local_guardrails_component_config(config_path: &str) -> serde_json::Map { - json!({ - "mode": "local", - "input": false, - "output": false, - "config_path": config_path, - "tool_input": true, - "tool_output": true, - "local": { - "python_module": "custom_guardrails" - } - }) - .as_object() - .unwrap() - .clone() -} - -fn local_llm_guardrails_component_config(config_yaml: &str) -> serde_json::Map { - json!({ - "mode": "local", - "codec": "openai_chat", - "input": true, - "output": true, - "config_yaml": config_yaml, - "colang_content": "define flow noop\n pass", - "local": { - "python_module": "custom_guardrails" - } - }) - .as_object() - .unwrap() - .clone() -} - #[test] fn target_scope_defaults_to_user_and_rejects_conflicts() { assert_eq!( @@ -472,64 +405,6 @@ fn typed_editor_model_contains_adaptive_options() { ); } -#[test] -fn typed_editor_model_contains_nemo_guardrails_options() { - let schema = NeMoGuardrailsConfig::editor_schema(); - assert!(!schema.fields.iter().any(|field| field.name == "version")); - assert_eq!( - schema.field("mode").unwrap().enum_values, - &["remote", "local"] - ); - assert_eq!(schema.field("codec").unwrap().kind, EditorFieldKind::Enum); - assert_eq!( - schema.field("input").unwrap().kind, - EditorFieldKind::Boolean - ); - assert_eq!( - schema.field("priority").unwrap().kind, - EditorFieldKind::Integer - ); - - let remote = schema.field("remote").unwrap().schema().unwrap(); - assert_eq!( - remote.field("timeout_millis").unwrap().kind, - EditorFieldKind::Integer - ); - assert_eq!( - remote.field("headers").unwrap().kind, - EditorFieldKind::StringMap - ); - - let local = schema.field("local").unwrap().schema().unwrap(); - assert_eq!( - local.field("python_module").unwrap().kind, - EditorFieldKind::String - ); - assert_eq!( - local.field("python_executable").unwrap().kind, - EditorFieldKind::String - ); - assert_eq!( - schema.field("config_path").unwrap().kind, - EditorFieldKind::String - ); - assert_eq!( - schema.field("config_yaml").unwrap().kind, - EditorFieldKind::String - ); - assert_eq!( - schema.field("colang_content").unwrap().kind, - EditorFieldKind::String - ); - - let request_defaults = schema.field("request_defaults").unwrap().schema().unwrap(); - let rails = request_defaults.field("rails").unwrap().schema().unwrap(); - assert_eq!( - rails.field("tool_input").unwrap().kind, - EditorFieldKind::Json - ); -} - #[test] fn typed_editor_model_contains_pii_redaction_options() { let schema = PiiRedactionConfig::editor_schema(); @@ -682,11 +557,6 @@ fn plugin_menu_builds_ordered_component_actions() { assert_eq!(items.len(), actions.len()); assert!(plain_labels[0].starts_with("Observability [on] —")); - assert!( - plain_labels - .iter() - .any(|label| { label.starts_with("NeMo Guardrails (Deprecated) [off] —") }) - ); assert_eq!( plain_labels[components.len()], "Example Dynamic — dynamic; config absent; schema fields" @@ -1253,63 +1123,6 @@ fn editor_model_adds_disabled_adaptive_component() { assert!(component.config.contains_key("policy")); } -#[test] -#[allow( - deprecated, - reason = "this compatibility test inspects the built-in Guardrails config until its removal" -)] -fn editor_model_reads_missing_nemo_guardrails_component_as_disabled_default() { - let config = PluginConfig::default(); - - let guardrails = component_nemo_guardrails_state(&config).unwrap(); - - assert!(!guardrails.enabled); - assert!( - !config - .components - .iter() - .any(|component| component.kind == NEMO_GUARDRAILS_PLUGIN_KIND) - ); - assert_eq!(guardrails.config.mode, "remote"); - assert!(!nemo_guardrails_configured(&guardrails.config)); - assert_eq!( - nemo_guardrails_summary(&guardrails), - "component disabled, fields none" - ); - assert!(!guardrails.should_store(nemo_guardrails_configured(&guardrails.config))); -} - -#[test] -fn editor_save_persists_disabled_nemo_guardrails_policy_only_edits() { - let mut config = PluginConfig::default(); - let mut guardrails = component_nemo_guardrails_state(&config).unwrap(); - let policy = NeMoGuardrailsConfig::editor_schema() - .field("policy") - .unwrap(); - - set_section_field( - &mut guardrails.config, - policy, - "unknown_field", - json!("ignore"), - ) - .unwrap(); - guardrails.mark_config_touched(); - - assert!(!guardrails.enabled); - assert!(!nemo_guardrails_configured(&guardrails.config)); - - store_nemo_guardrails_state(&mut config, &guardrails).unwrap(); - - let component = config - .components - .iter() - .find(|component| component.kind == NEMO_GUARDRAILS_PLUGIN_KIND) - .unwrap(); - assert!(!component.enabled); - assert_eq!(component.config["policy"]["unknown_field"], json!("ignore")); -} - #[test] fn typed_editor_serializes_explicit_observability_overrides() { let mut observability = ObservabilityConfig::default(); @@ -1553,93 +1366,6 @@ fn editor_save_preserves_unknown_adaptive_fields_and_all_sections() { ); } -#[test] -fn editor_save_preserves_unknown_nemo_guardrails_fields_and_sections() { - let mut config = PluginConfig { - components: vec![PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: true, - config: json!({ - "version": 1, - "future_top_level": "preserve", - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": "http://old.example.test", - "config_id": "old", - "future_remote": "preserve" - }, - "request_defaults": { - "future_defaults": "preserve", - "rails": { - "input": true, - "future_rails": "preserve" - } - } - }) - .as_object() - .unwrap() - .clone(), - }], - ..PluginConfig::default() - }; - let mut guardrails = component_nemo_guardrails_state(&config).unwrap(); - let schema = NeMoGuardrailsConfig::editor_schema(); - let remote = schema.field("remote").unwrap(); - let request_defaults = schema.field("request_defaults").unwrap(); - - set_struct_field(&mut guardrails.config, "codec", json!("openai_chat")).unwrap(); - set_section_field( - &mut guardrails.config, - remote, - "endpoint", - json!("http://localhost:8000"), - ) - .unwrap(); - set_section_field( - &mut guardrails.config, - remote, - "config_id", - json!("default"), - ) - .unwrap(); - set_section_field( - &mut guardrails.config, - request_defaults, - "context", - json!({"tenant": "docs"}), - ) - .unwrap(); - - guardrails.set_enabled(false); - store_nemo_guardrails_state(&mut config, &guardrails).unwrap(); - - let component = config - .components - .iter() - .find(|component| component.kind == NEMO_GUARDRAILS_PLUGIN_KIND) - .unwrap(); - assert!(!component.enabled); - assert!(!component.config.contains_key("version")); - assert_eq!( - component.config.get("future_top_level"), - Some(&json!("preserve")) - ); - let remote = component.config["remote"].as_object().unwrap(); - assert_eq!( - remote.get("endpoint"), - Some(&json!("http://localhost:8000")) - ); - assert_eq!(remote.get("future_remote"), Some(&json!("preserve"))); - let request_defaults = component.config["request_defaults"].as_object().unwrap(); - assert_eq!( - request_defaults.get("future_defaults"), - Some(&json!("preserve")) - ); - assert_eq!(request_defaults["context"], json!({"tenant": "docs"})); - assert_eq!(request_defaults["rails"]["future_rails"], json!("preserve")); -} - #[test] fn editor_save_preserves_unknown_pii_redaction_fields_and_prunes_version() { let mut config = PluginConfig { @@ -1985,55 +1711,6 @@ fn adaptive_summary_tracks_component_and_configured_fields() { ); } -#[test] -fn nemo_guardrails_summary_tracks_component_and_configured_fields() { - let config = PluginConfig::default(); - let mut guardrails = component_nemo_guardrails_state(&config).unwrap(); - - assert_eq!( - nemo_guardrails_summary(&guardrails), - "component disabled, fields none" - ); - - guardrails.set_enabled(true); - set_struct_field(&mut guardrails.config, "codec", json!("openai_chat")).unwrap(); - let remote = NeMoGuardrailsConfig::editor_schema() - .field("remote") - .unwrap(); - set_section_field( - &mut guardrails.config, - remote, - "endpoint", - json!("http://localhost:8000"), - ) - .unwrap(); - - assert!(nemo_guardrails_configured(&guardrails.config)); - assert_eq!( - nemo_guardrails_summary(&guardrails), - "component enabled, fields codec, remote" - ); - assert!(guardrails.should_store(nemo_guardrails_configured(&guardrails.config))); - - let existing = PluginConfig { - components: vec![PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: false, - config: guardrails_component_config("existing"), - }], - ..PluginConfig::default() - }; - let mut existing = component_nemo_guardrails_state(&existing).unwrap(); - reset_config_field( - &mut existing.config, - NeMoGuardrailsConfig::editor_schema() - .field("remote") - .unwrap(), - ) - .unwrap(); - assert!(existing.should_store(nemo_guardrails_configured(&existing.config))); -} - #[test] fn component_enablement_and_summary_track_config_state() { let mut config = PluginConfig::default(); @@ -2729,21 +2406,14 @@ fn write_plugin_config_prunes_defaults_and_round_trips() { enabled: true, config: adaptive_component_config("cli-roundtrip"), }); - config.components.push(PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: false, - config: guardrails_component_config("cli-roundtrip"), - }); - write_plugin_config(&path, &config).unwrap(); let rendered = std::fs::read_to_string(&path).unwrap(); assert!(rendered.contains("kind = \"observability\"")); assert!(rendered.contains("kind = \"adaptive\"")); - assert!(rendered.contains("kind = \"nemo_guardrails\"")); assert!(!rendered.contains("enabled = true")); let round_tripped = read_plugin_config(&path).unwrap(); - assert_eq!(round_tripped.components.len(), 3); + assert_eq!(round_tripped.components.len(), 2); assert_eq!(round_tripped.components[0].kind, OBSERVABILITY_PLUGIN_KIND); let adaptive = round_tripped .components @@ -2763,16 +2433,6 @@ fn write_plugin_config_prunes_defaults_and_round_trips() { adaptive_hints.get("inject_body_path"), Some(&json!("nvext.agent_hints")) ); - let guardrails = round_tripped - .components - .iter() - .find(|component| component.kind == NEMO_GUARDRAILS_PLUGIN_KIND) - .unwrap(); - assert!(!guardrails.enabled); - assert_eq!( - guardrails.config["remote"]["config_id"], - json!("cli-roundtrip") - ); } #[test] @@ -2877,34 +2537,6 @@ fn validate_config_accepts_adaptive_component() { validate_config(&config).unwrap(); } -#[test] -fn validate_config_accepts_nemo_guardrails_component() { - let config = PluginConfig { - components: vec![PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: true, - config: guardrails_component_config("cli-validation"), - }], - ..PluginConfig::default() - }; - - validate_config(&config).unwrap(); -} - -#[test] -fn validate_config_accepts_local_tool_only_nemo_guardrails_component() { - let config = PluginConfig { - components: vec![PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: true, - config: local_guardrails_component_config("./rails"), - }], - ..PluginConfig::default() - }; - - validate_config(&config).unwrap(); -} - #[test] fn validate_config_accepts_pii_redaction_component() { let config = PluginConfig { @@ -2931,213 +2563,6 @@ fn validate_config_accepts_pii_redaction_component() { validate_config(&config).unwrap(); } -#[test] -fn validate_config_rejects_local_nemo_guardrails_request_defaults() { - let config = PluginConfig { - components: vec![PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: true, - config: json!({ - "mode": "local", - "codec": "openai_chat", - "input": true, - "output": true, - "config_yaml": "models: []", - "request_defaults": { - "context": {"tenant": "demo"} - } - }) - .as_object() - .unwrap() - .clone(), - }], - ..PluginConfig::default() - }; - - let error = validate_config(&config).unwrap_err().to_string(); - assert!(error.contains("request_defaults"), "error was: {error}"); - assert!(error.contains("local mode"), "error was: {error}"); -} - -#[test] -fn validate_config_rejects_local_nemo_guardrails_multiple_config_sources() { - let config = PluginConfig { - components: vec![PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: true, - config: json!({ - "mode": "local", - "config_path": "./rails", - "config_yaml": "models: []" - }) - .as_object() - .unwrap() - .clone(), - }], - ..PluginConfig::default() - }; - - let error = validate_config(&config).unwrap_err().to_string(); - assert!( - error.contains("exactly one of config_path or config_yaml"), - "error was: {error}" - ); -} - -#[test] -fn validate_config_rejects_local_nemo_guardrails_colang_without_yaml() { - let config = PluginConfig { - components: vec![PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: true, - config: json!({ - "mode": "local", - "config_path": "./rails", - "colang_content": "define flow noop\n pass" - }) - .as_object() - .unwrap() - .clone(), - }], - ..PluginConfig::default() - }; - - let error = validate_config(&config).unwrap_err().to_string(); - assert!( - error.contains("colang_content can only be used with config_yaml"), - "error was: {error}" - ); -} - -#[test] -#[allow( - deprecated, - reason = "this compatibility test serializes the built-in Guardrails config until its removal" -)] -fn nemo_guardrails_config_map_prunes_default_version() { - let map = nemo_guardrails_config_map(&NeMoGuardrailsConfig { - codec: Some("openai_chat".into()), - remote: Some(RemoteBackendConfig { - endpoint: Some("http://localhost:8000".into()), - config_id: Some("default".into()), - ..RemoteBackendConfig::default() - }), - ..NeMoGuardrailsConfig::default() - }) - .unwrap(); - - assert!(!map.contains_key("version")); - assert_eq!(map.get("codec"), Some(&json!("openai_chat"))); - assert_eq!(map["remote"]["config_id"], json!("default")); -} - -#[test] -fn write_plugin_config_round_trips_local_nemo_guardrails_component() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("plugins.toml"); - let config = PluginConfig { - components: vec![PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: true, - config: local_guardrails_component_config("./rails"), - }], - ..PluginConfig::default() - }; - - write_plugin_config(&path, &config).unwrap(); - - let rendered = std::fs::read_to_string(&path).unwrap(); - assert!(rendered.contains("mode = \"local\"")); - assert!(rendered.contains("config_path = \"./rails\"")); - assert!(rendered.contains("tool_input = true")); - assert!(rendered.contains("python_module = \"custom_guardrails\"")); - - let round_tripped = read_plugin_config(&path).unwrap(); - let guardrails = round_tripped - .components - .iter() - .find(|component| component.kind == NEMO_GUARDRAILS_PLUGIN_KIND) - .unwrap(); - assert!(guardrails.enabled); - assert_eq!(guardrails.config["mode"], json!("local")); - assert_eq!(guardrails.config["config_path"], json!("./rails")); - assert_eq!(guardrails.config["tool_input"], json!(true)); - assert_eq!( - guardrails.config["local"]["python_module"], - json!("custom_guardrails") - ); -} - -#[test] -#[allow( - deprecated, - reason = "this compatibility test serializes the built-in Guardrails config until its removal" -)] -fn nemo_guardrails_config_map_serializes_local_mode_fields() { - let map = nemo_guardrails_config_map(&NeMoGuardrailsConfig { - mode: "local".into(), - config_path: Some("./rails".into()), - tool_input: true, - tool_output: true, - local: Some(LocalBackendConfig { - python_module: Some("custom_guardrails".into()), - python_executable: Some("/opt/python/bin/python3".into()), - python_path: None, - }), - ..NeMoGuardrailsConfig::default() - }) - .unwrap(); - - assert!(!map.contains_key("version")); - assert_eq!(map.get("mode"), Some(&json!("local"))); - assert_eq!(map.get("config_path"), Some(&json!("./rails"))); - assert_eq!(map.get("tool_input"), Some(&json!(true))); - assert_eq!(map["local"]["python_module"], json!("custom_guardrails")); - assert_eq!( - map["local"]["python_executable"], - json!("/opt/python/bin/python3") - ); -} - -#[test] -fn write_plugin_config_round_trips_local_llm_nemo_guardrails_component() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("plugins.toml"); - let config = PluginConfig { - components: vec![PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: true, - config: local_llm_guardrails_component_config("models: []"), - }], - ..PluginConfig::default() - }; - - write_plugin_config(&path, &config).unwrap(); - - let rendered = std::fs::read_to_string(&path).unwrap(); - assert!(rendered.contains("mode = \"local\"")); - assert!(rendered.contains("codec = \"openai_chat\"")); - assert!(rendered.contains("input = true")); - assert!(rendered.contains("output = true")); - assert!(rendered.contains("config_yaml = \"models: []\"")); - - let round_tripped = read_plugin_config(&path).unwrap(); - let guardrails = round_tripped - .components - .iter() - .find(|component| component.kind == NEMO_GUARDRAILS_PLUGIN_KIND) - .unwrap(); - assert_eq!(guardrails.config["mode"], json!("local")); - assert_eq!(guardrails.config["codec"], json!("openai_chat")); - assert_eq!(guardrails.config["input"], json!(true)); - assert_eq!(guardrails.config["output"], json!(true)); - assert_eq!(guardrails.config["config_yaml"], json!("models: []")); - assert_eq!( - guardrails.config["colang_content"], - json!("define flow noop\n pass") - ); -} - #[test] fn display_helpers_render_scalars_json_and_defaults() { assert_eq!(display_value(&json!("logs")), "logs"); diff --git a/crates/cli/tests/coverage/shared/server_tests.rs b/crates/cli/tests/coverage/shared/server_tests.rs index 7628f122b3..abd2ca8136 100644 --- a/crates/cli/tests/coverage/shared/server_tests.rs +++ b/crates/cli/tests/coverage/shared/server_tests.rs @@ -2676,6 +2676,28 @@ fn register_and_validate_plugin_components_rejects_legacy_switchyard_components( } } +#[test] +fn register_and_validate_plugin_components_rejects_legacy_nemo_guardrails_components() { + for enabled in [true, false] { + let config = PluginConfig { + components: vec![PluginComponentSpec { + kind: "nemo_guardrails".into(), + enabled, + config: Map::new(), + }], + ..PluginConfig::default() + }; + + let errors = register_and_validate_plugin_components(&config); + assert!( + errors + .iter() + .any(|error| matches!(error, PluginComponentSetupError::RemovedNemoGuardrails)), + "legacy NeMo Guardrails components must be rejected when enabled is {enabled}" + ); + } +} + #[test] fn plugin_component_setup_errors_render_every_diagnostic_variant() { let adaptive = PluginComponentSetupError::Adaptive("adaptive failure".into()); @@ -2706,6 +2728,16 @@ fn plugin_component_setup_errors_render_every_diagnostic_variant() { .contains("removed in NeMo Relay >=0.8.0") ); assert!(switchyard.to_string().contains("migration guide")); + + let guardrails = PluginComponentSetupError::RemovedNemoGuardrails; + assert_eq!(guardrails.check_name(), "NeMo Guardrails migration"); + assert_eq!(guardrails.diagnostic_details(), guardrails.to_string()); + assert!( + guardrails + .to_string() + .contains("removed in NeMo Relay >=0.10.0") + ); + assert!(guardrails.to_string().contains("migration guide")); } fn dynamic_component_without_manifest( @@ -2755,6 +2787,20 @@ async fn plugin_host_activation_covers_empty_invalid_and_missing_manifest_paths( assert!(dynamic_switchyard.contains("has no manifest_ref")); assert!(!dynamic_switchyard.contains("removed in NeMo Relay 0.8")); + let dynamic_guardrails = activate_server_plugins( + None, + vec![dynamic_component_without_manifest( + "nemo_guardrails", + DynamicPluginKind::Worker, + )], + ) + .await + .err() + .expect("dynamic NeMo Guardrails plugin without a manifest should reach dynamic activation"); + let dynamic_guardrails = dynamic_guardrails.to_string(); + assert!(dynamic_guardrails.contains("has no manifest_ref")); + assert!(!dynamic_guardrails.contains("removed in NeMo Relay 0.10")); + let worker = activate_server_plugins( None, vec![dynamic_component_without_manifest( diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index 462ba99156..4061158a1d 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -15,7 +15,6 @@ workspace = true [features] default = [ - "guardrails-remote", "object-store", ] atof-streaming = [ @@ -30,8 +29,6 @@ schema = ["dep:schemars", "nemo-relay-types/schema"] __skip-implicit-config = [] # Private compatibility-free test harness for migrating legacy lifecycle tests. __test-plugin-host = [] -# Deprecated with the built-in NeMo Guardrails integration; scheduled for removal in 0.9. -guardrails-remote = [] object-store = [ "dep:object_store", "tokio/net", diff --git a/crates/core/src/plugin.rs b/crates/core/src/plugin.rs index e12ee02667..eea2b135f3 100644 --- a/crates/core/src/plugin.rs +++ b/crates/core/src/plugin.rs @@ -135,6 +135,9 @@ pub(crate) const ATIF_RUNTIME_DELIVERY_FAILURE_MARKER: &str = "ATIF runtime deli pub(crate) const OTEL_RUNTIME_DELIVERY_FAILURE_MARKER: &str = "OpenTelemetry runtime delivery failures"; +const REMOVED_NEMO_GUARDRAILS_COMPONENT_KIND: &str = "nemo_guardrails"; +const REMOVED_NEMO_GUARDRAILS_COMPONENT_MESSAGE: &str = "the built-in NeMo Guardrails integration was removed in NeMo Relay >=0.10.0; remove this `[[components]]` entry and refer to the migration guide: https://docs.nvidia.com/nemo/relay/reference/migration-guides#remove-the-built-in-nemo-guardrails-component"; + /// Canonical plugin configuration document. #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] @@ -1165,10 +1168,6 @@ fn register_plugin_with_owner( /// /// Built-in plugins are available to validation and initialization without a /// binding or application-specific registration call. -#[allow( - deprecated, - reason = "the host must register the built-in Guardrails plugin until its scheduled removal" -)] pub fn ensure_builtin_plugins_registered() -> Result<()> { let all_registered = { let guard = PLUGIN_HANDLERS.read().map_err(|err| { @@ -1176,7 +1175,6 @@ pub fn ensure_builtin_plugins_registered() -> Result<()> { })?; [ crate::observability::plugin_component::OBSERVABILITY_PLUGIN_KIND, - crate::plugins::nemo_guardrails::component::NEMO_GUARDRAILS_PLUGIN_KIND, crate::plugins::model_pricing::PRICING_PLUGIN_KIND, crate::plugins::resource_metrics::RESOURCE_METRICS_PLUGIN_KIND, ] @@ -1195,7 +1193,6 @@ pub fn ensure_builtin_plugins_registered() -> Result<()> { // call so a removed built-in is restored, a replacement is rejected, and // a corrected ownership conflict can be retried without restarting Relay. crate::observability::plugin_component::register_observability_component()?; - crate::plugins::nemo_guardrails::component::register_nemo_guardrails_component()?; crate::plugins::model_pricing::register_pricing_component()?; crate::plugins::resource_metrics::register_resource_metrics_component() } @@ -1323,6 +1320,13 @@ fn lookup_registered_plugin(plugin_kind: &str) -> Option> { /// kinds, and plugin-provided validation hooks. #[doc(hidden)] pub fn validate_static_plugin_config(config: &PluginConfig) -> ConfigReport { + validate_plugin_config_with_static_component_count(config, config.components.len()) +} + +fn validate_plugin_config_with_static_component_count( + config: &PluginConfig, + static_component_count: usize, +) -> ConfigReport { let mut report = ConfigReport::default(); if let Err(error) = ensure_builtin_plugins_registered() { report.diagnostics.push(ConfigDiagnostic { @@ -1348,7 +1352,19 @@ pub fn validate_static_plugin_config(config: &PluginConfig) -> ConfigReport { validate_plugin_multiplicity(&mut report, config); - for component in &config.components { + for (index, component) in config.components.iter().enumerate() { + if index < static_component_count + && component.kind == REMOVED_NEMO_GUARDRAILS_COMPONENT_KIND + { + report.diagnostics.push(ConfigDiagnostic { + level: DiagnosticLevel::Error, + code: "plugin.removed_component".to_string(), + component: Some(component.kind.clone()), + field: None, + message: REMOVED_NEMO_GUARDRAILS_COMPONENT_MESSAGE.to_string(), + }); + continue; + } let Some(plugin) = lookup_registered_plugin(&component.kind) else { push_policy_diag( &mut report.diagnostics, @@ -1641,16 +1657,24 @@ fn plugin_mutation_executor() -> Result<&'static PluginMutationSender> { pub(crate) async fn initialize_plugins_exact_for_host( config: PluginConfig, + static_component_count: usize, owner_id: u64, rollback_failures: Arc>>, diagnostics: Vec, ) -> Result { verify_plugin_host_owner(owner_id)?; - initialize_plugins_exact_inner(config, Some(rollback_failures), diagnostics).await + initialize_plugins_exact_inner( + config, + static_component_count, + Some(rollback_failures), + diagnostics, + ) + .await } async fn initialize_plugins_exact_inner( config: PluginConfig, + static_component_count: usize, rollback_failures: Option>>>, diagnostics: Vec, ) -> Result { @@ -1669,9 +1693,10 @@ async fn initialize_plugins_exact_inner( diagnostics, ..ConfigReport::default() }; - report - .diagnostics - .extend(validate_static_plugin_config(&config).diagnostics); + report.diagnostics.extend( + validate_plugin_config_with_static_component_count(&config, static_component_count) + .diagnostics, + ); if report.has_errors() { return Err(PluginError::InvalidConfig(join_error_messages(&report))); } diff --git a/crates/core/src/plugin/dynamic/host.rs b/crates/core/src/plugin/dynamic/host.rs index ca9755e50e..759125c013 100644 --- a/crates/core/src/plugin/dynamic/host.rs +++ b/crates/core/src/plugin/dynamic/host.rs @@ -163,6 +163,7 @@ impl PluginHostActivation { dynamic_plugins: Vec, diagnostics: Vec, ) -> Result<(Self, ConfigReport)> { + let static_component_count = config.components.len(); let dynamic_plugin_count = dynamic_plugins.len(); log::info!( target: "nemo_relay.plugin", @@ -238,6 +239,7 @@ impl PluginHostActivation { let owner_id = claim.owner_id(); let initialization = tokio::spawn(initialize_plugins_exact_for_host( config, + static_component_count, owner_id, Arc::clone(&rollback_failures), diagnostics, diff --git a/crates/core/src/plugins/mod.rs b/crates/core/src/plugins/mod.rs index 75a5c8af09..6f206aedd4 100644 --- a/crates/core/src/plugins/mod.rs +++ b/crates/core/src/plugins/mod.rs @@ -4,5 +4,4 @@ //! First-party plugin implementations for NeMo Relay Core. pub mod model_pricing; -pub mod nemo_guardrails; pub mod resource_metrics; diff --git a/crates/core/src/plugins/nemo_guardrails/component.rs b/crates/core/src/plugins/nemo_guardrails/component.rs deleted file mode 100644 index 3505d2d633..0000000000 --- a/crates/core/src/plugins/nemo_guardrails/component.rs +++ /dev/null @@ -1,1445 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! NeMo Guardrails plugin component contract. -#![allow( - deprecated, - reason = "the built-in implementation remains supported until its scheduled removal" -)] - -use std::collections::HashMap; -use std::future::Future; -use std::pin::Pin; -use std::sync::Arc; - -use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value as Json}; - -use crate::codec::resolve::supported_codec_names; -use crate::plugin::{ - ConfigDiagnostic, ConfigPolicy, DiagnosticLevel, Plugin, PluginComponentSpec, PluginError, - PluginRegistrationContext, Result as PluginResult, UnsupportedBehavior, - apply_global_config_policy, deregister_plugin, register_builtin_plugin, -}; - -#[path = "local.rs"] -mod local; -#[cfg(feature = "guardrails-remote")] -#[path = "remote.rs"] -mod remote; -use local::register_local_backend; -#[cfg(feature = "guardrails-remote")] -use remote::register_remote_backend; - -/// The built-in NeMo Guardrails plugin kind. -#[deprecated( - since = "0.8.0", - note = "the built-in NeMo Guardrails plugin is scheduled for removal in 0.9; no replacement is available in 0.8" -)] -pub const NEMO_GUARDRAILS_PLUGIN_KIND: &str = "nemo_guardrails"; - -/// Stable diagnostic code for the built-in plugin's deprecation warning. -const NEMO_GUARDRAILS_DEPRECATION_CODE: &str = "nemo_guardrails.deprecated"; - -/// NeMo Relay release in which the built-in plugin is scheduled for removal. -const NEMO_GUARDRAILS_REMOVAL_VERSION: &str = "0.9"; - -const NEMO_GUARDRAILS_DEPRECATION_MESSAGE: &str = "the built-in `nemo_guardrails` plugin is deprecated and scheduled for removal in NeMo Relay 0.9"; - -#[cfg(not(feature = "guardrails-remote"))] -fn register_remote_backend( - _config: NeMoGuardrailsConfig, - _ctx: &mut PluginRegistrationContext, -) -> PluginResult<()> { - Err(PluginError::RegistrationFailed( - "built-in NeMo Guardrails remote backend is unavailable in this build".to_string(), - )) -} - -/// Top-level NeMo Guardrails component wrapper. -#[derive(Debug, Clone)] -#[deprecated( - since = "0.8.0", - note = "the built-in NeMo Guardrails plugin is scheduled for removal in 0.9; no replacement is available in 0.8" -)] -pub struct ComponentSpec { - /// Whether the component should be activated. - pub enabled: bool, - /// Component-local NeMo Guardrails config. - pub config: NeMoGuardrailsConfig, -} - -impl ComponentSpec { - /// Creates an enabled NeMo Guardrails component spec. - pub fn new(config: NeMoGuardrailsConfig) -> Self { - Self { - enabled: true, - config, - } - } -} - -impl From for PluginComponentSpec { - fn from(value: ComponentSpec) -> Self { - let Json::Object(config) = serde_json::to_value(value.config) - .expect("NeMo Guardrails config should serialize to an object") - else { - unreachable!("NeMo Guardrails config must serialize to an object"); - }; - - PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: value.enabled, - config, - } - } -} - -/// Canonical config document for the deprecated built-in NeMo Guardrails component. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[deprecated( - since = "0.8.0", - note = "the built-in NeMo Guardrails plugin is scheduled for removal in 0.9; no replacement is available in 0.8" -)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -pub struct NeMoGuardrailsConfig { - /// NeMo Guardrails config schema version. - #[serde(default = "default_nemo_guardrails_config_version")] - pub version: u32, - /// Backend mode: `remote` or `local`. - #[serde(default = "default_mode")] - #[cfg_attr(feature = "schema", schemars(schema_with = "mode_schema"))] - pub mode: String, - /// Path to a native NeMo Guardrails config directory. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub config_path: Option, - /// Inline native NeMo Guardrails YAML config. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub config_yaml: Option, - /// Optional inline Colang content. Valid only with `config_yaml`. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub colang_content: Option, - /// Provider request/response codec for LLM-managed surfaces. - #[serde(default, skip_serializing_if = "Option::is_none")] - #[cfg_attr(feature = "schema", schemars(schema_with = "codec_schema"))] - pub codec: Option, - /// Whether to run input rails around managed LLM execution. - #[serde(default = "default_true")] - pub input: bool, - /// Whether to run output rails around managed LLM execution. - #[serde(default = "default_true")] - pub output: bool, - /// Whether to run tool-input rails around managed tool execution. - #[serde(default)] - pub tool_input: bool, - /// Whether to run tool-output rails around managed tool execution. - #[serde(default)] - pub tool_output: bool, - /// Intercept priority. Lower values run earlier. - #[serde(default = "default_priority")] - pub priority: i32, - /// Remote-backend settings used when `mode = "remote"`. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub remote: Option, - /// Local-backend settings used when `mode = "local"`. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub local: Option, - /// Default request semantics passed through to the selected Guardrails backend. - /// - /// This models request-time concepts such as rail selection and generation - /// options without claiming backend parity for every Guardrails feature. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub request_defaults: Option, - /// Component-local unsupported-config policy. - #[serde(default)] - pub policy: ConfigPolicy, -} - -impl Default for NeMoGuardrailsConfig { - fn default() -> Self { - Self { - version: default_nemo_guardrails_config_version(), - mode: default_mode(), - config_path: None, - config_yaml: None, - colang_content: None, - codec: None, - input: true, - output: true, - tool_input: false, - tool_output: false, - priority: default_priority(), - remote: None, - local: None, - request_defaults: None, - policy: ConfigPolicy::default(), - } - } -} - -/// Remote-backend settings for a hosted NeMo Guardrails service. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[deprecated( - since = "0.8.0", - note = "the built-in NeMo Guardrails plugin is scheduled for removal in 0.9; no replacement is available in 0.8" -)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -pub struct RemoteBackendConfig { - /// Base URL for the remote Guardrails service. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub endpoint: Option, - /// One remote Guardrails config identifier. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub config_id: Option, - /// Multiple remote Guardrails config identifiers to combine. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub config_ids: Vec, - /// Static request headers sent to the remote service. - #[serde(default)] - pub headers: HashMap, - /// Request timeout in milliseconds. - #[serde(default = "default_timeout_millis")] - pub timeout_millis: u64, -} - -impl Default for RemoteBackendConfig { - fn default() -> Self { - Self { - endpoint: None, - config_id: None, - config_ids: vec![], - headers: HashMap::new(), - timeout_millis: default_timeout_millis(), - } - } -} - -/// Local-backend settings for the Python `nemoguardrails` runtime. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[deprecated( - since = "0.8.0", - note = "the built-in NeMo Guardrails plugin is scheduled for removal in 0.9; no replacement is available in 0.8" -)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -pub struct LocalBackendConfig { - /// Optional import path for the Python runtime module. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub python_module: Option, - /// Optional Python executable used to run the local Guardrails worker. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub python_executable: Option, - /// Optional PYTHONPATH used only by the local Guardrails worker subprocess. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub python_path: Option, -} - -/// Default request semantics applied by the selected Guardrails backend. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[deprecated( - since = "0.8.0", - note = "the built-in NeMo Guardrails plugin is scheduled for removal in 0.9; no replacement is available in 0.8" -)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -pub struct RequestDefaultsConfig { - /// Default context object passed into Guardrails requests. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub context: Option, - /// Default remote thread identifier for continuation-aware requests. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub thread_id: Option, - /// Default remote Guardrails state payload for continuation-aware requests. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub state: Option, - /// Default request-time rail selection. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub rails: Option, - /// Default model parameters applied to Guardrails-backed LLM calls. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub llm_params: Option, - /// Whether to include raw LLM output in Guardrails responses. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub llm_output: Option, - /// Default output variables selection. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub output_vars: Option, - /// Default generation-log selection. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub log: Option, -} - -/// Request-time rail selection for Guardrails generation. -/// -/// These are backend request options, not top-level NeMo Relay interception -/// surfaces. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[deprecated( - since = "0.8.0", - note = "the built-in NeMo Guardrails plugin is scheduled for removal in 0.9; no replacement is available in 0.8" -)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -pub struct RequestRailsConfig { - /// Input rails selection. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub input: Option, - /// Output rails selection. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub output: Option, - /// Retrieval rails selection. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub retrieval: Option, - /// Dialog rails selection. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub dialog: Option, - /// Tool-output rails selection. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tool_output: Option, - /// Tool-input rails selection. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tool_input: Option, -} - -/// Rail-selection shape used by Guardrails generation options. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(untagged)] -#[deprecated( - since = "0.8.0", - note = "the built-in NeMo Guardrails plugin is scheduled for removal in 0.9; no replacement is available in 0.8" -)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -pub enum RailSelector { - /// Enable or disable the whole rail family. - Enabled(bool), - /// Enable only named rails within a family. - Named(Vec), -} - -crate::editor_config! { - impl NeMoGuardrailsConfig { - mode => { - label: "mode", - kind: Enum, - values: ["remote", "local"], - }, - config_path => { label: "config_path", kind: String, optional: true }, - config_yaml => { label: "config_yaml", kind: String, optional: true }, - colang_content => { label: "colang_content", kind: String, optional: true }, - codec => { - label: "codec", - kind: Enum, - values: ["openai_chat", "openai_responses", "anthropic_messages", "oci_genai", "gemini_generate_content"], - optional: true, - }, - input => { label: "input", kind: Boolean }, - output => { label: "output", kind: Boolean }, - tool_input => { label: "tool_input", kind: Boolean }, - tool_output => { label: "tool_output", kind: Boolean }, - priority => { label: "priority", kind: Integer }, - remote => { - label: "remote", - kind: Section, - optional: true, - nested: RemoteBackendConfig, - default: RemoteBackendConfig, - }, - local => { - label: "local", - kind: Section, - optional: true, - nested: LocalBackendConfig, - default: LocalBackendConfig, - }, - request_defaults => { - label: "request_defaults", - kind: Section, - optional: true, - nested: RequestDefaultsConfig, - default: RequestDefaultsConfig, - }, - policy => { - label: "policy", - kind: Section, - nested: ConfigPolicy, - default: ConfigPolicy, - }, - } -} - -crate::editor_config! { - impl RemoteBackendConfig { - endpoint => { label: "endpoint", kind: String, optional: true }, - config_id => { label: "config_id", kind: String, optional: true }, - config_ids => { label: "config_ids", kind: List, list: &crate::config_editor::STRING_LIST_ITEM }, - headers => { label: "headers", kind: StringMap }, - timeout_millis => { label: "timeout_millis", kind: Integer }, - } -} - -crate::editor_config! { - impl LocalBackendConfig { - python_module => { label: "python_module", kind: String, optional: true }, - python_executable => { label: "python_executable", kind: String, optional: true }, - python_path => { label: "python_path", kind: String, optional: true }, - } -} - -crate::editor_config! { - impl RequestDefaultsConfig { - context => { label: "context", kind: Json, optional: true }, - thread_id => { label: "thread_id", kind: String, optional: true }, - state => { label: "state", kind: Json, optional: true }, - rails => { - label: "rails", - kind: Section, - optional: true, - nested: RequestRailsConfig, - default: RequestRailsConfig, - }, - llm_params => { label: "llm_params", kind: Json, optional: true }, - llm_output => { label: "llm_output", kind: Boolean, optional: true }, - output_vars => { label: "output_vars", kind: Json, optional: true }, - log => { label: "log", kind: Json, optional: true }, - } -} - -crate::editor_config! { - impl RequestRailsConfig { - input => { label: "input", kind: Json, optional: true }, - output => { label: "output", kind: Json, optional: true }, - retrieval => { label: "retrieval", kind: Json, optional: true }, - dialog => { label: "dialog", kind: Boolean, optional: true }, - tool_output => { label: "tool_output", kind: Json, optional: true }, - tool_input => { label: "tool_input", kind: Json, optional: true }, - } -} - -struct NeMoGuardrailsPlugin; - -impl Plugin for NeMoGuardrailsPlugin { - fn plugin_kind(&self) -> &str { - NEMO_GUARDRAILS_PLUGIN_KIND - } - - fn allows_multiple_components(&self) -> bool { - false - } - - fn validate(&self, plugin_config: &Map) -> Vec { - validate_nemo_guardrails_plugin_config(plugin_config) - } - - fn validate_with_policy( - &self, - plugin_config: &Map, - policy: &ConfigPolicy, - ) -> Vec { - validate_nemo_guardrails_plugin_config_with_policy(plugin_config, Some(policy)) - } - - fn register<'a>( - &'a self, - plugin_config: &Map, - ctx: &'a mut PluginRegistrationContext, - ) -> Pin> + Send + 'a>> { - let parsed = parse_nemo_guardrails_config(plugin_config); - Box::pin(async move { - let config = parsed?; - register_nemo_guardrails_backend(config, ctx) - }) - } -} - -/// Registers the `nemo_guardrails` component kind in the plugin registry. -#[deprecated( - since = "0.8.0", - note = "the built-in NeMo Guardrails plugin is scheduled for removal in 0.9; no replacement is available in 0.8" -)] -pub fn register_nemo_guardrails_component() -> PluginResult<()> { - register_builtin_plugin(Arc::new(NeMoGuardrailsPlugin)) -} - -/// Deregisters the `nemo_guardrails` component kind from the plugin registry. -#[deprecated( - since = "0.8.0", - note = "the built-in NeMo Guardrails plugin is scheduled for removal in 0.9; no replacement is available in 0.8" -)] -pub fn deregister_nemo_guardrails_component() -> bool { - deregister_plugin(NEMO_GUARDRAILS_PLUGIN_KIND) -} - -/// Returns the JSON Schema for the NeMo Guardrails component configuration. -#[cfg(feature = "schema")] -#[deprecated( - since = "0.8.0", - note = "the built-in NeMo Guardrails plugin is scheduled for removal in 0.9; no replacement is available in 0.8" -)] -pub fn nemo_guardrails_config_schema() -> serde_json::Value { - serde_json::to_value(schemars::schema_for!(NeMoGuardrailsConfig)) - .expect("NeMo Guardrails config schema should serialize") -} - -#[cfg(feature = "schema")] -fn mode_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::schema::Schema { - string_enum_schema(generator, &["remote", "local"], Some("remote")) -} - -#[cfg(feature = "schema")] -fn codec_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::schema::Schema { - string_enum_schema( - generator, - &[ - "openai_chat", - "openai_responses", - "anthropic_messages", - "oci_genai", - "gemini_generate_content", - ], - None, - ) -} - -#[cfg(feature = "schema")] -fn string_enum_schema( - generator: &mut schemars::r#gen::SchemaGenerator, - values: &[&str], - default: Option<&str>, -) -> schemars::schema::Schema { - let mut schema: schemars::schema::SchemaObject = - ::json_schema(generator).into(); - schema.enum_values = Some( - values - .iter() - .map(|value| Json::String((*value).into())) - .collect(), - ); - if let Some(default) = default { - schema.metadata().default = Some(Json::String(default.into())); - } - schema.into() -} - -fn register_nemo_guardrails_backend( - config: NeMoGuardrailsConfig, - ctx: &mut PluginRegistrationContext, -) -> PluginResult<()> { - log::warn!( - target: "nemo_relay.plugin", - event = "nemo_guardrails_deprecated", - plugin_kind = NEMO_GUARDRAILS_PLUGIN_KIND, - removal_version = NEMO_GUARDRAILS_REMOVAL_VERSION; - "The built-in NeMo Guardrails plugin is deprecated and scheduled for removal in NeMo Relay 0.9" - ); - - match config.mode.as_str() { - "remote" => register_remote_backend(config, ctx), - "local" => register_local_backend(config, ctx), - other => Err(PluginError::InvalidConfig(format!( - "unsupported NeMo Guardrails mode '{other}'" - ))), - } -} - -fn parse_nemo_guardrails_config( - plugin_config: &Map, -) -> PluginResult { - serde_json::from_value(Json::Object(plugin_config.clone())).map_err(|err| { - PluginError::InvalidConfig(format!("invalid NeMo Guardrails plugin config: {err}")) - }) -} - -fn validate_nemo_guardrails_plugin_config( - plugin_config: &Map, -) -> Vec { - validate_nemo_guardrails_plugin_config_with_policy(plugin_config, None) -} - -fn validate_nemo_guardrails_plugin_config_with_policy( - plugin_config: &Map, - policy: Option<&ConfigPolicy>, -) -> Vec { - let deprecation = nemo_guardrails_deprecation_diagnostic(); - let mut config = match parse_nemo_guardrails_config(plugin_config) { - Ok(config) => config, - Err(err) => { - return vec![ - deprecation, - ConfigDiagnostic { - level: DiagnosticLevel::Error, - code: "nemo_guardrails.invalid_plugin_config".to_string(), - component: Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - field: None, - message: err.to_string(), - }, - ]; - } - }; - if let Some(policy) = policy { - config.policy = apply_global_config_policy(config.policy, policy); - } - - let mut diagnostics = vec![deprecation]; - - validate_unknown_fields( - &mut diagnostics, - &config.policy, - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - plugin_config, - &[ - "version", - "mode", - "config_path", - "config_yaml", - "colang_content", - "codec", - "input", - "output", - "tool_input", - "tool_output", - "priority", - "remote", - "local", - "request_defaults", - "policy", - ], - ); - - validate_policy_fields(&mut diagnostics, &config.policy, plugin_config); - validate_section_fields( - &mut diagnostics, - &config.policy, - plugin_config, - "remote", - &[ - "endpoint", - "config_id", - "config_ids", - "headers", - "timeout_millis", - ], - ); - validate_section_fields( - &mut diagnostics, - &config.policy, - plugin_config, - "local", - &["python_module", "python_executable", "python_path"], - ); - validate_section_fields( - &mut diagnostics, - &config.policy, - plugin_config, - "request_defaults", - &[ - "context", - "thread_id", - "state", - "rails", - "llm_params", - "llm_output", - "output_vars", - "log", - ], - ); - validate_nested_section_fields( - &mut diagnostics, - &config.policy, - plugin_config, - "request_defaults", - "rails", - &[ - "input", - "output", - "retrieval", - "dialog", - "tool_output", - "tool_input", - ], - ); - - validate_version(&mut diagnostics, &config.policy, config.version); - validate_mode(&mut diagnostics, &config.policy, &config.mode); - validate_non_empty_strings(&mut diagnostics, &config.policy, &config); - validate_config_shape(&mut diagnostics, &config.policy, &config); - validate_codec_requirements(&mut diagnostics, &config.policy, &config); - validate_surface_selection(&mut diagnostics, &config.policy, &config); - validate_remote_backend_support(&mut diagnostics, &config.policy, &config); - validate_request_defaults(&mut diagnostics, &config.policy, &config); - - diagnostics -} - -fn nemo_guardrails_deprecation_diagnostic() -> ConfigDiagnostic { - ConfigDiagnostic { - level: DiagnosticLevel::Warning, - code: NEMO_GUARDRAILS_DEPRECATION_CODE.to_string(), - component: Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - field: None, - message: NEMO_GUARDRAILS_DEPRECATION_MESSAGE.to_string(), - } -} - -fn validate_version(diagnostics: &mut Vec, policy: &ConfigPolicy, version: u32) { - if version != 1 { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_config_version", - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - Some("version".to_string()), - format!("NeMo Guardrails config version {version} is unsupported"), - ); - } -} - -fn validate_mode(diagnostics: &mut Vec, policy: &ConfigPolicy, mode: &str) { - if !matches!(mode, "remote" | "local") { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - Some("mode".to_string()), - "mode must be 'remote' or 'local'".to_string(), - ); - } -} - -fn validate_non_empty_strings( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - config: &NeMoGuardrailsConfig, -) { - validate_optional_non_empty_string( - diagnostics, - policy, - "config_path", - config.config_path.as_deref(), - "config_path must not be empty", - ); - validate_optional_non_empty_string( - diagnostics, - policy, - "config_yaml", - config.config_yaml.as_deref(), - "config_yaml must not be empty", - ); - validate_optional_non_empty_string( - diagnostics, - policy, - "colang_content", - config.colang_content.as_deref(), - "colang_content must not be empty", - ); - - if let Some(remote) = &config.remote { - validate_remote_non_empty_strings(diagnostics, policy, remote); - } - - if let Some(local) = &config.local { - validate_local_non_empty_strings(diagnostics, policy, local); - } -} - -fn validate_remote_non_empty_strings( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - remote: &RemoteBackendConfig, -) { - validate_optional_non_empty_string( - diagnostics, - policy, - "remote.endpoint", - remote.endpoint.as_deref(), - "remote.endpoint must not be empty", - ); - validate_optional_non_empty_string( - diagnostics, - policy, - "remote.config_id", - remote.config_id.as_deref(), - "remote.config_id must not be empty", - ); - for (index, config_id) in remote.config_ids.iter().enumerate() { - validate_optional_non_empty_string( - diagnostics, - policy, - format!("remote.config_ids[{index}]"), - Some(config_id.as_str()), - "remote.config_ids entries must not be empty", - ); - } -} - -fn validate_local_non_empty_strings( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - local: &LocalBackendConfig, -) { - validate_optional_non_empty_string( - diagnostics, - policy, - "local.python_module", - local.python_module.as_deref(), - "local.python_module must not be empty", - ); - validate_optional_non_empty_string( - diagnostics, - policy, - "local.python_executable", - local.python_executable.as_deref(), - "local.python_executable must not be empty", - ); - validate_optional_non_empty_string( - diagnostics, - policy, - "local.python_path", - local.python_path.as_deref(), - "local.python_path must not be empty", - ); -} - -fn validate_optional_non_empty_string( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - field: impl Into, - value: Option<&str>, - message: &str, -) { - if let Some(value) = value - && value.trim().is_empty() - { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - Some(field.into()), - message.to_string(), - ); - } -} - -fn validate_config_shape( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - config: &NeMoGuardrailsConfig, -) { - let flags = ConfigShapeFlags::from(config); - - match config.mode.as_str() { - "local" => validate_local_config_shape(diagnostics, policy, config, &flags), - "remote" => validate_remote_config_shape(diagnostics, policy, config, &flags), - _ => {} - } -} - -struct ConfigShapeFlags { - has_config_path: bool, - has_config_yaml: bool, - has_colang_content: bool, - has_remote_config_id: bool, - has_remote_config_ids: bool, -} - -impl From<&NeMoGuardrailsConfig> for ConfigShapeFlags { - fn from(config: &NeMoGuardrailsConfig) -> Self { - Self { - has_config_path: config.config_path.is_some(), - has_config_yaml: config.config_yaml.is_some(), - has_colang_content: config.colang_content.is_some(), - has_remote_config_id: config - .remote - .as_ref() - .and_then(|remote| remote.config_id.as_ref()) - .is_some(), - has_remote_config_ids: config - .remote - .as_ref() - .map(|remote| !remote.config_ids.is_empty()) - .unwrap_or(false), - } - } -} - -fn validate_local_config_shape( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - config: &NeMoGuardrailsConfig, - flags: &ConfigShapeFlags, -) { - if flags.has_config_path == flags.has_config_yaml { - push_config_shape_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.invalid_config_source", - None, - "exactly one of config_path or config_yaml is required in local mode", - ); - } - - if flags.has_colang_content && !flags.has_config_yaml { - push_config_shape_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some("colang_content"), - "colang_content can only be used with config_yaml", - ); - } - - if config.remote.is_some() { - push_config_shape_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some("remote"), - "remote backend settings cannot be used when mode is 'local'", - ); - } -} - -fn validate_remote_config_shape( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - config: &NeMoGuardrailsConfig, - flags: &ConfigShapeFlags, -) { - if flags.has_config_path || flags.has_config_yaml || flags.has_colang_content { - push_config_shape_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.invalid_config_source", - None, - "remote mode uses remote config identity and cannot include config_path, config_yaml, or colang_content", - ); - } - - if config.local.is_some() { - push_config_shape_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some("local"), - "local backend settings cannot be used when mode is 'remote'", - ); - } - - match &config.remote { - Some(remote) - if remote - .endpoint - .as_ref() - .is_some_and(|value| !value.trim().is_empty()) => {} - _ => push_config_shape_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some("remote.endpoint"), - "remote.endpoint is required when mode is 'remote'", - ), - } - - if flags.has_remote_config_id && flags.has_remote_config_ids { - push_config_shape_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some("remote"), - "remote.config_id and remote.config_ids cannot be used together", - ); - } - - if !(flags.has_remote_config_id || flags.has_remote_config_ids) { - push_config_shape_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.invalid_config_source", - None, - "remote mode requires remote.config_id or remote.config_ids", - ); - } -} - -fn push_config_shape_diag( - diagnostics: &mut Vec, - behavior: UnsupportedBehavior, - code: &str, - field: Option<&str>, - message: &str, -) { - push_policy_diag( - diagnostics, - behavior, - code, - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - field.map(str::to_string), - message.to_string(), - ); -} - -fn validate_codec_requirements( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - config: &NeMoGuardrailsConfig, -) { - let llm_surface_enabled = config.input || config.output; - if !llm_surface_enabled { - return; - } - - let Some(codec) = config.codec.as_deref() else { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - Some("codec".to_string()), - "codec is required when any LLM surface is enabled".to_string(), - ); - return; - }; - - if !supported_codec_names().contains(&codec) { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - Some("codec".to_string()), - format!( - "codec must be one of: {}", - supported_codec_names().join(", ") - ), - ); - } -} - -fn validate_surface_selection( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - config: &NeMoGuardrailsConfig, -) { - if config.input || config.output || config.tool_input || config.tool_output { - return; - } - - push_policy_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - None, - "at least one Guardrails surface must be enabled".to_string(), - ); -} - -fn validate_remote_backend_support( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - config: &NeMoGuardrailsConfig, -) { - if config.mode != "remote" { - return; - } - - if (config.input || config.output) - && config - .codec - .as_deref() - .is_some_and(|codec| codec != "openai_chat") - { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - Some("codec".to_string()), - "remote mode currently supports only codec = 'openai_chat'".to_string(), - ); - } - - if config.tool_input { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - Some("tool_input".to_string()), - "remote mode does not currently support managed tool_input against the stock Guardrails remote contract".to_string(), - ); - } -} - -fn validate_request_defaults( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - config: &NeMoGuardrailsConfig, -) { - let Some(request_defaults) = &config.request_defaults else { - return; - }; - - if config.mode == "local" { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - Some("request_defaults".to_string()), - "local mode does not currently support request_defaults".to_string(), - ); - return; - } - - validate_json_object_field( - diagnostics, - policy, - request_defaults.context.as_ref(), - "request_defaults.context", - "request_defaults.context must be a JSON object", - ); - validate_request_thread_id(diagnostics, policy, request_defaults.thread_id.as_deref()); - validate_json_object_field( - diagnostics, - policy, - request_defaults.state.as_ref(), - "request_defaults.state", - "request_defaults.state must be a JSON object", - ); - validate_request_state_keys(diagnostics, policy, request_defaults.state.as_ref()); - validate_json_object_field( - diagnostics, - policy, - request_defaults.llm_params.as_ref(), - "request_defaults.llm_params", - "request_defaults.llm_params must be a JSON object", - ); - validate_json_object_field( - diagnostics, - policy, - request_defaults.log.as_ref(), - "request_defaults.log", - "request_defaults.log must be a JSON object", - ); - - validate_output_vars(diagnostics, policy, request_defaults.output_vars.as_ref()); - validate_request_rails(diagnostics, policy, request_defaults.rails.as_ref()); -} - -fn push_request_defaults_diag( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - field: &str, - message: &str, -) { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - Some(field.to_string()), - message.to_string(), - ); -} - -fn validate_request_thread_id( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - thread_id: Option<&str>, -) { - let Some(thread_id) = thread_id else { - return; - }; - - let trimmed_thread_id = thread_id.trim(); - if trimmed_thread_id.is_empty() { - push_request_defaults_diag( - diagnostics, - policy, - "request_defaults.thread_id", - "request_defaults.thread_id must not be empty", - ); - } else if trimmed_thread_id.len() < 16 { - push_request_defaults_diag( - diagnostics, - policy, - "request_defaults.thread_id", - "request_defaults.thread_id must be at least 16 characters long", - ); - } -} - -fn validate_request_state_keys( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - state: Option<&Json>, -) { - let Some(state) = state.and_then(Json::as_object) else { - return; - }; - - let contains_supported_key = state.contains_key("events") || state.contains_key("state"); - let contains_unsupported_key = state.keys().any(|key| key != "events" && key != "state"); - if (!state.is_empty() && !contains_supported_key) || contains_unsupported_key { - push_request_defaults_diag( - diagnostics, - policy, - "request_defaults.state", - "request_defaults.state must be empty or contain only 'events' or 'state'", - ); - } -} - -fn validate_output_vars( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - output_vars: Option<&Json>, -) { - let Some(output_vars) = output_vars else { - return; - }; - - match output_vars { - Json::Bool(_) => {} - Json::Array(values) => validate_output_var_entries(diagnostics, policy, values), - _ => push_request_defaults_diag( - diagnostics, - policy, - "request_defaults.output_vars", - "request_defaults.output_vars must be a boolean or an array of strings", - ), - } -} - -fn validate_output_var_entries( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - values: &[Json], -) { - for (index, value) in values.iter().enumerate() { - if !value.is_string() || value.as_str().is_some_and(|entry| entry.trim().is_empty()) { - push_request_defaults_diag( - diagnostics, - policy, - &format!("request_defaults.output_vars[{index}]"), - "request_defaults.output_vars array entries must be non-empty strings", - ); - } - } -} - -fn validate_request_rails( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - rails: Option<&RequestRailsConfig>, -) { - let Some(rails) = rails else { - return; - }; - - validate_rail_selector( - diagnostics, - policy, - rails.input.as_ref(), - "request_defaults.rails.input", - ); - validate_rail_selector( - diagnostics, - policy, - rails.output.as_ref(), - "request_defaults.rails.output", - ); - validate_rail_selector( - diagnostics, - policy, - rails.retrieval.as_ref(), - "request_defaults.rails.retrieval", - ); - validate_rail_selector( - diagnostics, - policy, - rails.tool_output.as_ref(), - "request_defaults.rails.tool_output", - ); - validate_rail_selector( - diagnostics, - policy, - rails.tool_input.as_ref(), - "request_defaults.rails.tool_input", - ); -} - -fn validate_json_object_field( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - value: Option<&Json>, - field: &str, - message: &str, -) { - let Some(value) = value else { - return; - }; - - if !value.is_object() { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - Some(field.to_string()), - message.to_string(), - ); - } -} - -fn validate_rail_selector( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - value: Option<&RailSelector>, - field: &str, -) { - let Some(value) = value else { - return; - }; - - if let RailSelector::Named(names) = value { - for (index, name) in names.iter().enumerate() { - if name.trim().is_empty() { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "nemo_guardrails.unsupported_value", - Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), - Some(format!("{field}[{index}]")), - "named rail selections must not contain empty strings".to_string(), - ); - } - } - } -} - -fn validate_policy_fields( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - plugin_config: &Map, -) { - if let Some(policy_json) = plugin_config.get("policy").and_then(Json::as_object) { - validate_unknown_fields( - diagnostics, - policy, - Some("policy".to_string()), - policy_json, - &["unknown_component", "unknown_field", "unsupported_value"], - ); - } -} - -fn validate_section_fields( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - plugin_config: &Map, - section: &str, - known_fields: &[&str], -) { - if let Some(section_json) = plugin_config.get(section).and_then(Json::as_object) { - validate_unknown_fields( - diagnostics, - policy, - Some(section.to_string()), - section_json, - known_fields, - ); - } -} - -fn validate_nested_section_fields( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - plugin_config: &Map, - section: &str, - nested_section: &str, - known_fields: &[&str], -) { - if let Some(section_json) = plugin_config.get(section).and_then(Json::as_object) - && let Some(nested_json) = section_json.get(nested_section).and_then(Json::as_object) - { - validate_unknown_fields( - diagnostics, - policy, - Some(format!("{section}.{nested_section}")), - nested_json, - known_fields, - ); - } -} - -fn validate_unknown_fields( - diagnostics: &mut Vec, - policy: &ConfigPolicy, - component: Option, - config: &Map, - known_fields: &[&str], -) { - for field in config.keys() { - if !known_fields.contains(&field.as_str()) { - push_policy_diag( - diagnostics, - policy.unknown_field, - "nemo_guardrails.unknown_field", - component.clone(), - Some(field.clone()), - format!( - "field '{}' is not recognized for '{}'", - field, - component.as_deref().unwrap_or("unknown") - ), - ); - } - } -} - -fn push_policy_diag( - diagnostics: &mut Vec, - behavior: UnsupportedBehavior, - code: &str, - component: Option, - field: Option, - message: String, -) { - let level = match behavior { - UnsupportedBehavior::Ignore => return, - UnsupportedBehavior::Warn => DiagnosticLevel::Warning, - UnsupportedBehavior::Error => DiagnosticLevel::Error, - }; - - diagnostics.push(ConfigDiagnostic { - level, - code: code.to_string(), - component, - field, - message, - }); -} - -fn default_nemo_guardrails_config_version() -> u32 { - 1 -} - -fn default_mode() -> String { - "remote".to_string() -} - -fn default_true() -> bool { - true -} - -fn default_priority() -> i32 { - 100 -} - -fn default_timeout_millis() -> u64 { - 3_000 -} - -#[cfg(test)] -#[path = "../../../tests/unit/plugins/nemo_guardrails/component_tests.rs"] -mod tests; diff --git a/crates/core/src/plugins/nemo_guardrails/local.rs b/crates/core/src/plugins/nemo_guardrails/local.rs deleted file mode 100644 index e1618836f1..0000000000 --- a/crates/core/src/plugins/nemo_guardrails/local.rs +++ /dev/null @@ -1,15 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use crate::plugin::{PluginRegistrationContext, Result as PluginResult}; - -use super::NeMoGuardrailsConfig; - -mod python; - -pub(super) fn register_local_backend( - config: NeMoGuardrailsConfig, - ctx: &mut PluginRegistrationContext, -) -> PluginResult<()> { - python::register_local_backend(config, ctx) -} diff --git a/crates/core/src/plugins/nemo_guardrails/local_worker.py b/crates/core/src/plugins/nemo_guardrails/local_worker.py deleted file mode 100644 index 6bb54022cc..0000000000 --- a/crates/core/src/plugins/nemo_guardrails/local_worker.py +++ /dev/null @@ -1,311 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import asyncio -import importlib -import json -import sys -import traceback -from collections.abc import AsyncIterator, MutableSet -from typing import Any, TypeAlias - -DEFAULT_MODULE_NAME = "nemoguardrails" -SUPPORTED_NEMOGUARDRAILS_VERSION = "0.22.0" -STREAM_QUEUE_MAXSIZE = 32 - -JsonObject: TypeAlias = dict[str, Any] -StreamQueue: TypeAlias = asyncio.Queue[str | None] -PendingTasks: TypeAlias = MutableSet[asyncio.Task[None]] - -_PROTOCOL_STDOUT = sys.stdout -sys.stdout = sys.stderr - - -def send(message: JsonObject) -> None: - _PROTOCOL_STDOUT.write(json.dumps(message, separators=(",", ":")) + "\n") - _PROTOCOL_STDOUT.flush() - - -def response(request_id: str, result: Any | None = None) -> None: - payload = {"id": request_id, "ok": True} - if result is not None: - payload["result"] = result - send(payload) - - -def error_response(request_id: str, error: BaseException) -> None: - send({"id": request_id, "ok": False, "error": str(error)}) - - -def stream_event(request_id: str, event: str, **fields: Any) -> None: - payload = {"id": request_id, "ok": True, "event": event} - payload.update(fields) - send(payload) - - -def stream_error(request_id: str, error: BaseException) -> None: - send({"id": request_id, "ok": False, "event": "error", "error": str(error)}) - - -def status_value(status: Any) -> str: - value = getattr(status, "value", status) - return str(value).lower() - - -def optional_string_attr(obj: Any, attr: str) -> str | None: - value = getattr(obj, attr, None) - if value is None: - return None - return str(value) - - -def string_attr_or_empty(obj: Any, attr: str) -> str: - return optional_string_attr(obj, attr) or "" - - -def guardrails_stream_error_message(chunk: str) -> str | None: - try: - payload = json.loads(chunk) - except Exception: - return None - error = payload.get("error") - if not isinstance(error, dict): - return None - if error.get("type") != "guardrails_violation": - return None - return error.get("message") or "Blocked by output rails." - - -class AsyncTextStream: - def __init__(self, queue: StreamQueue) -> None: - self._queue = queue - - def __aiter__(self) -> AsyncIterator[str]: - return self - - async def __anext__(self) -> str: - value = await self._queue.get() - if value is None: - raise StopAsyncIteration - return value - - -class GuardrailsWorker: - def __init__(self, config: JsonObject) -> None: - if sys.version_info < (3, 11): - raise RuntimeError("NeMo Guardrails local backend requires python3 >= 3.11") - - local = config.get("local") or {} - root_module = (local.get("python_module") or DEFAULT_MODULE_NAME).strip() - guardrails = self._import_dependency(root_module, root_module) - options = self._import_dependency(f"{root_module}.rails.llm.options", root_module) - - version = getattr(guardrails, "__version__", None) - if version != SUPPORTED_NEMOGUARDRAILS_VERSION: - raise RuntimeError( - "NeMo Guardrails local backend requires " - f"nemoguardrails=={SUPPORTED_NEMOGUARDRAILS_VERSION}, but found {version!r}. " - f"Install it with: pip install nemoguardrails=={SUPPORTED_NEMOGUARDRAILS_VERSION}" - ) - - self._rail_type = options.RailType - self._rail_status = options.RailStatus - guardrails_config = self._build_guardrails_config(guardrails.RailsConfig, config) - self._rails = guardrails.LLMRails(guardrails_config) - - def _import_dependency(self, module_name: str, root_module: str) -> Any: - try: - return importlib.import_module(module_name) - except ImportError as err: - missing = getattr(err, "name", None) - if missing == root_module: - raise RuntimeError( - "NeMo Guardrails is required for the built-in NeMo Guardrails local backend. " - f"Install it with: pip install nemoguardrails=={SUPPORTED_NEMOGUARDRAILS_VERSION}" - ) from err - raise RuntimeError( - "NeMo Guardrails local backend could not import a required dependency: " - f"{missing or err}. Install the full NeMo Guardrails runtime dependencies." - ) from err - - def _build_guardrails_config(self, rails_config_cls: Any, config: JsonObject) -> Any: - config_path = config.get("config_path") - if config_path: - return rails_config_cls.from_path(config_path) - - config_yaml = config.get("config_yaml") - if config_yaml is None: - raise ValueError("config_yaml is required when config_path is not provided") - return rails_config_cls.from_content( - colang_content=config.get("colang_content"), - yaml_content=config_yaml, - ) - - def _rail_kind(self, rail_type: str | None) -> Any: - if rail_type == "input": - return self._rail_type.INPUT - if rail_type == "output": - return self._rail_type.OUTPUT - raise ValueError(f"unsupported rail_type {rail_type!r}") - - async def check(self, messages: list[JsonObject], rail_type: str | None) -> JsonObject: - result = await self._rails.check_async( - messages, - rail_types=[self._rail_kind(rail_type)], - ) - return { - "status": status_value(result.status), - "content": string_attr_or_empty(result, "content"), - "rail": optional_string_attr(result, "rail"), - } - - def has_streaming_output_rails(self) -> bool: - output = self._output_rails_config() - flows = getattr(output, "flows", None) if output is not None else None - return bool(flows) - - def ensure_streaming_output_supported(self) -> None: - output = self._output_rails_config() - if output is None: - return - - streaming = getattr(output, "streaming", None) - if streaming is None or not bool(getattr(streaming, "enabled", False)): - raise RuntimeError( - "local NeMo Guardrails streaming output rails require " - "rails.output.streaming.enabled = true in the Guardrails config." - ) - - if not bool(getattr(streaming, "stream_first", True)): - raise RuntimeError( - "local NeMo Guardrails streaming output rails currently require " - "rails.output.streaming.stream_first = true." - ) - - def _output_rails_config(self) -> Any: - config = getattr(self._rails, "config", None) - rails = getattr(config, "rails", None) - return getattr(rails, "output", None) - - async def monitor_stream( - self, request_id: str, messages: list[JsonObject], queue: StreamQueue, streams: dict[str, StreamQueue] - ) -> None: - try: - async for chunk in self._rails.stream_async( - messages=messages, - generator=AsyncTextStream(queue), - include_metadata=False, - ): - if not isinstance(chunk, str): - continue - message = guardrails_stream_error_message(chunk) - if message: - stream_event(request_id, "blocked", message=message) - return - stream_event(request_id, "done") - except Exception as err: - stream_error(request_id, err) - finally: - streams.pop(request_id, None) - - -worker: GuardrailsWorker | None = None -streams: dict[str, StreamQueue] = {} - - -def track_task(pending_tasks: PendingTasks, task: asyncio.Task[None]) -> asyncio.Task[None]: - pending_tasks.add(task) - task.add_done_callback(pending_tasks.discard) - return task - - -async def handle_message(message: JsonObject, pending_tasks: PendingTasks) -> None: - global worker - - request_id = str(message.get("id", "")) - command = message.get("command") - try: - if command == "init": - worker = _initialize_worker(message) - response(request_id, _worker_details()) - elif worker is None: - raise RuntimeError("NeMo Guardrails local Python worker is not initialized") - else: - await _handle_worker_command(worker, command, request_id, message, pending_tasks) - except Exception as err: - if command and command.startswith("stream_"): - stream_error(request_id, err) - else: - error_response(request_id, err) - - -def _initialize_worker(message: JsonObject) -> GuardrailsWorker: - return GuardrailsWorker(message.get("config") or {}) - - -def _worker_details() -> JsonObject: - return {"python": sys.executable, "version": ".".join(str(part) for part in sys.version_info[:3])} - - -async def _handle_worker_command( - worker: GuardrailsWorker, - command: str | None, - request_id: str, - message: JsonObject, - pending_tasks: PendingTasks, -) -> None: - if command == "check": - response(request_id, await worker.check(message.get("messages") or [], message.get("rail_type"))) - elif command == "has_streaming_output_rails": - response(request_id, {"enabled": worker.has_streaming_output_rails()}) - elif command == "ensure_streaming_output_supported": - worker.ensure_streaming_output_supported() - response(request_id) - elif command == "stream_start": - _start_stream(worker, request_id, message, pending_tasks) - elif command in {"stream_text", "stream_end"}: - await _write_stream(command, request_id, message) - else: - raise RuntimeError(f"unknown worker command {command!r}") - - -def _start_stream(worker: GuardrailsWorker, request_id: str, message: JsonObject, pending_tasks: PendingTasks) -> None: - queue = asyncio.Queue(maxsize=STREAM_QUEUE_MAXSIZE) - streams[request_id] = queue - task = worker.monitor_stream(request_id, message.get("messages") or [], queue, streams) - track_task(pending_tasks, asyncio.create_task(task)) - - -async def _write_stream(command: str | None, request_id: str, message: JsonObject) -> None: - queue = streams.get(request_id) - if queue is not None: - await queue.put(message.get("text") or "" if command == "stream_text" else None) - - -async def main() -> None: - pending_tasks = set() - try: - while True: - line = await asyncio.to_thread(sys.stdin.readline) - if not line: - return - try: - message = json.loads(line) - except Exception: - traceback.print_exc(file=sys.stderr) - continue - if str(message.get("command", "")).startswith("stream_"): - await handle_message(message, pending_tasks) - else: - track_task( - pending_tasks, - asyncio.create_task(handle_message(message, pending_tasks)), - ) - finally: - for task in tuple(pending_tasks): - task.cancel() - if pending_tasks: - await asyncio.gather(*pending_tasks, return_exceptions=True) - - -asyncio.run(main()) diff --git a/crates/core/src/plugins/nemo_guardrails/mod.rs b/crates/core/src/plugins/nemo_guardrails/mod.rs deleted file mode 100644 index 5909d9a28b..0000000000 --- a/crates/core/src/plugins/nemo_guardrails/mod.rs +++ /dev/null @@ -1,14 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Deprecated built-in NeMo Guardrails plugin integration for NeMo Relay Core. - -#[cfg(test)] -use std::sync::Mutex; - -#[cfg(test)] -pub(crate) fn test_mutex() -> &'static Mutex<()> { - crate::shared_runtime::runtime_owner_test_mutex() -} - -pub mod component; diff --git a/crates/core/src/plugins/nemo_guardrails/python.rs b/crates/core/src/plugins/nemo_guardrails/python.rs deleted file mode 100644 index 11a4180f48..0000000000 --- a/crates/core/src/plugins/nemo_guardrails/python.rs +++ /dev/null @@ -1,1488 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use std::collections::HashMap; -use std::env; -use std::ffi::{OsStr, OsString}; -use std::future::Future; -use std::io::{BufRead, BufReader, Write}; -use std::pin::Pin; -use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio}; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use std::sync::{Arc, Mutex, mpsc as std_mpsc}; -use std::task::{Context, Poll}; -use std::thread; -use std::time::Duration; - -use serde::Deserialize; -use serde_json::json; -use tokio::sync::{mpsc, watch}; -use tokio::task::JoinHandle; -use tokio_stream::StreamExt; -use tokio_stream::wrappers::ReceiverStream; - -use crate::api::llm::LlmRequest; -use crate::api::runtime::{ - LlmExecutionFn, LlmJsonStream, LlmStreamExecutionFn, LlmStreamInner, ToolExecutionFn, -}; -use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; -use crate::codec::resolve::{ProviderSurface, request_codec, response_codec}; -use crate::error::{FlowError, Result as FlowResult}; -use crate::json::Json; -use crate::plugin::{PluginError, PluginRegistrationContext, Result as PluginResult}; - -use super::NeMoGuardrailsConfig; - -#[cfg(not(windows))] -const DEFAULT_PYTHON_EXECUTABLE: &str = "python3"; -#[cfg(windows)] -const DEFAULT_PYTHON_EXECUTABLE: &str = "python"; -const PYTHON_EXECUTABLE_ENV: &str = "NEMO_RELAY_PYTHON"; -const PYO3_PYTHON_ENV: &str = "PYO3_PYTHON"; -const UV_PYTHON_ENV: &str = "UV_PYTHON"; -const WORKER_INIT_TIMEOUT: Duration = Duration::from_secs(30); -const WORKER_RPC_TIMEOUT: Duration = Duration::from_secs(30); -const WORKER_SCRIPT: &str = include_str!("local_worker.py"); - -pub(super) fn register_local_backend( - config: NeMoGuardrailsConfig, - ctx: &mut PluginRegistrationContext, -) -> PluginResult<()> { - let runtime = Arc::new(LocalGuardrailsRuntime::new(&config)?); - - if config.input || config.output { - let llm_runtime = Arc::clone(&runtime); - let enable_input = config.input; - let enable_output = config.output; - let llm_execution: LlmExecutionFn = Arc::new(move |_name, request, _context, next| { - let runtime = Arc::clone(&llm_runtime); - Box::pin(async move { - runtime - .execute_llm(request, next, enable_input, enable_output) - .await - }) - }); - ctx.register_llm_execution_intercept( - "nemo_guardrails_local", - config.priority, - llm_execution, - )?; - - let stream_runtime = Arc::clone(&runtime); - let enable_input = config.input; - let enable_output = config.output; - let llm_stream_execution: LlmStreamExecutionFn = - Arc::new(move |_name, request, _context, next| { - let runtime = Arc::clone(&stream_runtime); - Box::pin(async move { - runtime - .execute_llm_stream(request, next, enable_input, enable_output) - .await - }) - }); - ctx.register_llm_stream_execution_intercept( - "nemo_guardrails_local_stream", - config.priority, - llm_stream_execution, - )?; - } - - if config.tool_input || config.tool_output { - let tool_runtime = Arc::clone(&runtime); - let enable_tool_input = config.tool_input; - let enable_tool_output = config.tool_output; - let tool_execution: ToolExecutionFn = Arc::new(move |context, next| { - let runtime = Arc::clone(&tool_runtime); - let tool_name = context.tool_name().to_string(); - let args = context.into_args(); - Box::pin(async move { - let current_args = if enable_tool_input { - runtime.check_tool_input(&tool_name, &args).await? - } else { - args - }; - - let mut execution_result = next(current_args.clone()).await?; - if enable_tool_output { - execution_result.result = runtime - .check_tool_output(&tool_name, ¤t_args, &execution_result.result) - .await?; - } - Ok(execution_result.into()) - }) - }); - ctx.register_tool_execution_intercept( - "nemo_guardrails_local", - config.priority, - tool_execution, - )?; - } - - Ok(()) -} - -struct LocalGuardrailsRuntime { - bridge: LocalGuardrailsBridge, - codec: Option, -} - -impl LocalGuardrailsRuntime { - fn new(config: &NeMoGuardrailsConfig) -> PluginResult { - Ok(Self { - bridge: LocalGuardrailsBridge::new(config)?, - codec: resolve_codec(config)?, - }) - } - - async fn execute_llm( - &self, - request: LlmRequest, - next: crate::api::runtime::LlmExecutionNextFn, - enable_input: bool, - enable_output: bool, - ) -> FlowResult { - let (request, messages) = self.prepare_llm_request(request, enable_input).await?; - let response = next(request).await?; - - if enable_output { - let annotated_response = self.codec()?.decode_response(&response)?; - if let Some(response_text) = annotated_response.response_text() { - self.check_output_rails(&messages, response_text).await?; - } - } - - Ok(response) - } - - async fn execute_llm_stream( - &self, - request: LlmRequest, - next: crate::api::runtime::LlmStreamExecutionNextFn, - enable_input: bool, - enable_output: bool, - ) -> FlowResult { - let (request, messages) = self.prepare_llm_request(request, enable_input).await?; - let provider_stream = next(request).await?; - - if !enable_output || !self.bridge.has_streaming_output_rails().await? { - return Ok(provider_stream); - } - - self.bridge.ensure_streaming_output_supported().await?; - self.guard_provider_stream(messages, provider_stream).await - } - - async fn prepare_llm_request( - &self, - request: LlmRequest, - enable_input: bool, - ) -> FlowResult<(LlmRequest, Vec)> { - let codec = self.codec()?; - let mut current_request = request; - let mut annotated = codec.decode(¤t_request)?; - let mut messages = messages_from_annotated(&annotated)?; - - if enable_input { - match self - .bridge - .check(messages.clone(), LocalRailKind::Input) - .await? - { - LocalCheckOutcome::Passed => {} - LocalCheckOutcome::Blocked { rail, .. } => { - return Err(blocked_error("input", rail.as_deref())); - } - LocalCheckOutcome::Modified { content, .. } => { - replace_last_role_content(&mut annotated, "user", content)?; - current_request = codec.encode(&annotated, ¤t_request)?; - messages = messages_from_annotated(&annotated)?; - } - } - } - - Ok((current_request, messages)) - } - - async fn check_output_rails(&self, messages: &[Json], response_text: &str) -> FlowResult<()> { - let mut output_messages = messages.to_vec(); - output_messages.push(json!({ - "role": "assistant", - "content": response_text, - })); - - match self - .bridge - .check(output_messages, LocalRailKind::Output) - .await? - { - LocalCheckOutcome::Passed => Ok(()), - LocalCheckOutcome::Blocked { rail, .. } => { - Err(blocked_error("output", rail.as_deref())) - } - LocalCheckOutcome::Modified { .. } => Err(local_violation( - "NeMo Guardrails output rail returned modified content, but the local backend \ - does not rewrite provider responses yet.", - )), - } - } - - async fn check_tool_input(&self, tool_name: &str, args: &Json) -> FlowResult { - let messages = vec![json!({ - "role": "user", - "content": tool_input_content(tool_name, args)?, - })]; - - match self.bridge.check(messages, LocalRailKind::Input).await? { - LocalCheckOutcome::Passed => Ok(args.clone()), - LocalCheckOutcome::Blocked { rail, .. } => { - Err(blocked_error("tool_input", rail.as_deref())) - } - LocalCheckOutcome::Modified { content, .. } => { - modified_tool_payload(&content, "arguments") - } - } - } - - async fn check_tool_output( - &self, - tool_name: &str, - args: &Json, - result: &Json, - ) -> FlowResult { - let messages = vec![ - json!({ - "role": "user", - "content": tool_input_content(tool_name, args)?, - }), - json!({ - "role": "assistant", - "content": tool_output_content(tool_name, args, result)?, - }), - ]; - - match self.bridge.check(messages, LocalRailKind::Output).await? { - LocalCheckOutcome::Passed => Ok(result.clone()), - LocalCheckOutcome::Blocked { rail, .. } => { - Err(blocked_error("tool_output", rail.as_deref())) - } - LocalCheckOutcome::Modified { content, .. } => { - modified_tool_payload(&content, "result") - } - } - } - - async fn guard_provider_stream( - &self, - messages: Vec, - provider_stream: LlmJsonStream, - ) -> FlowResult { - let (text_tx, text_rx) = mpsc::channel::>(32); - let (chunk_tx, chunk_rx) = mpsc::channel::>(32); - let blocked = Arc::new(Mutex::new(None)); - let monitor = self - .bridge - .spawn_stream_monitor(messages, text_rx, Arc::clone(&blocked))?; - let codec = *self.codec()?; - - let (cancel, cancel_rx) = watch::channel(false); - let (closed, closed_rx) = watch::channel(None); - tokio::spawn(async move { - forward_guarded_provider_stream( - provider_stream, - codec, - text_tx, - chunk_tx, - monitor, - blocked, - cancel_rx, - closed, - ) - .await; - }); - - Ok(LlmJsonStream::from_closeable(GuardedProviderStream { - receiver: ReceiverStream::new(chunk_rx), - cancel, - closed: closed_rx, - })) - } - - fn codec(&self) -> FlowResult<&LocalGuardrailsCodec> { - self.codec.as_ref().ok_or_else(|| { - FlowError::Internal( - "local NeMo Guardrails backend requires a supported codec".to_string(), - ) - }) - } -} - -struct LocalGuardrailsBridge { - worker: Arc, -} - -impl LocalGuardrailsBridge { - fn new(config: &NeMoGuardrailsConfig) -> PluginResult { - Ok(Self { - worker: LocalGuardrailsWorker::start(config)?, - }) - } - - async fn check( - &self, - messages: Vec, - kind: LocalRailKind, - ) -> FlowResult { - let result = self - .worker - .request(json!({ - "command": "check", - "messages": messages, - "rail_type": kind.as_str(), - })) - .await?; - parse_check_result(result) - } - - async fn has_streaming_output_rails(&self) -> FlowResult { - let result = self - .worker - .request(json!({ "command": "has_streaming_output_rails" })) - .await?; - result - .get("enabled") - .and_then(Json::as_bool) - .ok_or_else(|| FlowError::Internal("worker returned invalid streaming probe".into())) - } - - async fn ensure_streaming_output_supported(&self) -> FlowResult<()> { - self.worker - .request(json!({ "command": "ensure_streaming_output_supported" })) - .await - .map(|_| ()) - } - - fn spawn_stream_monitor( - &self, - messages: Vec, - text_rx: mpsc::Receiver>, - blocked: Arc>>, - ) -> FlowResult>> { - let (stream_id, event_rx) = self.worker.start_stream(messages)?; - let worker = Arc::clone(&self.worker); - Ok(tokio::spawn(async move { - monitor_guardrails_stream(worker, stream_id, text_rx, event_rx, blocked).await - })) - } -} - -struct LocalGuardrailsWorker { - writer: Mutex>, - child: Mutex, - waiters: Arc>>>, - stream_events: Arc>>>, - next_id: AtomicU64, - shutdown_started: AtomicBool, -} - -impl LocalGuardrailsWorker { - fn start(config: &NeMoGuardrailsConfig) -> PluginResult> { - log::info!( - target: "nemo_relay.worker", - event = "worker_starting", - plugin_id = "nemo_guardrails"; - "NeMo Guardrails local worker is starting" - ); - let python = python_executable(config); - let mut command = Command::new(&python); - command - .arg("-u") - .arg("-c") - .arg(WORKER_SCRIPT) - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::inherit()); - if let Some(python_path) = worker_python_path(config) { - command.env("PYTHONPATH", python_path); - } - - let mut child = command.spawn().map_err(|err| { - PluginError::RegistrationFailed(format!( - "failed to start NeMo Guardrails local Python worker with {python:?}: {err}" - )) - })?; - let stdin = child.stdin.take().ok_or_else(|| { - PluginError::RegistrationFailed( - "failed to open stdin for NeMo Guardrails local Python worker".to_string(), - ) - })?; - let stdout = child.stdout.take().ok_or_else(|| { - PluginError::RegistrationFailed( - "failed to open stdout for NeMo Guardrails local Python worker".to_string(), - ) - })?; - - let worker = Arc::new(Self { - writer: Mutex::new(Some(WorkerCommandWriter::spawn(stdin))), - child: Mutex::new(child), - waiters: Arc::new(Mutex::new(HashMap::new())), - stream_events: Arc::new(Mutex::new(HashMap::new())), - next_id: AtomicU64::new(1), - shutdown_started: AtomicBool::new(false), - }); - worker.spawn_reader(stdout); - worker.initialize(config)?; - log::info!( - target: "nemo_relay.plugin", - event = "plugin_resource_access_validated", - plugin_kind = "nemo_guardrails", - resource_kind = "python_worker", - permission = "execute"; - "Plugin resource access validated" - ); - log::info!( - target: "nemo_relay.worker", - event = "worker_connected", - plugin_id = "nemo_guardrails"; - "NeMo Guardrails local worker connected" - ); - Ok(worker) - } - - fn spawn_reader(&self, stdout: ChildStdout) { - let waiters = Arc::clone(&self.waiters); - let stream_events = Arc::clone(&self.stream_events); - thread::spawn(move || { - let reader = BufReader::new(stdout); - for line in reader.lines() { - let line = match line { - Ok(line) => line, - Err(err) => { - notify_worker_closed(&waiters, &stream_events, err.to_string()); - return; - } - }; - if line.trim().is_empty() { - continue; - } - let envelope = match serde_json::from_str::(&line) { - Ok(envelope) => envelope, - Err(err) => { - notify_worker_closed( - &waiters, - &stream_events, - format!("invalid worker response: {err}"), - ); - return; - } - }; - dispatch_worker_envelope(&waiters, &stream_events, envelope); - } - notify_worker_closed(&waiters, &stream_events, "worker exited".to_string()); - }); - } - - fn initialize(&self, config: &NeMoGuardrailsConfig) -> PluginResult<()> { - let response = self - .request_blocking( - json!({ - "command": "init", - "config": config, - }), - WORKER_INIT_TIMEOUT, - ) - .map_err(|err| PluginError::RegistrationFailed(err.to_string()))?; - if response.ok { - Ok(()) - } else { - Err(PluginError::RegistrationFailed( - response - .error - .unwrap_or_else(|| "NeMo Guardrails local Python worker failed".to_string()), - )) - } - } - - async fn request(&self, payload: Json) -> FlowResult { - self.request_with_timeout(payload, WORKER_RPC_TIMEOUT).await - } - - async fn request_with_timeout(&self, mut payload: Json, timeout: Duration) -> FlowResult { - let receiver = self.send_request(&mut payload)?; - let response_task = tokio::task::spawn_blocking(move || receiver.recv()); - let envelope = match tokio::time::timeout(timeout, response_task).await { - Ok(result) => result - .map_err(|err| FlowError::Internal(format!("worker response task failed: {err}")))? - .map_err(|err| { - FlowError::Internal(format!("worker response channel closed: {err}")) - })?, - Err(_) => { - log::error!( - target: "nemo_relay.worker", - event = "worker_failed", - plugin_id = "nemo_guardrails", - reason = "request_timeout"; - "NeMo Guardrails local worker request timed out" - ); - self.shutdown(); - return Err(FlowError::Internal(format!( - "worker request timed out after {} seconds", - timeout.as_secs() - ))); - } - }; - worker_result(envelope) - } - - fn request_blocking(&self, mut payload: Json, timeout: Duration) -> FlowResult { - let receiver = self.send_request(&mut payload)?; - receiver - .recv_timeout(timeout) - .map_err(|err| FlowError::Internal(format!("worker did not initialize: {err}"))) - } - - fn send_request(&self, payload: &mut Json) -> FlowResult> { - let id = self.next_request_id(); - set_request_id(payload, &id)?; - let (tx, rx) = std_mpsc::channel(); - self.waiters - .lock() - .map_err(|err| FlowError::Internal(format!("worker waiter lock poisoned: {err}")))? - .insert(id.clone(), tx); - if let Err(err) = self.write_command(payload) { - let _ = self.waiters.lock().map(|mut waiters| waiters.remove(&id)); - return Err(err); - } - Ok(rx) - } - - fn start_stream( - &self, - messages: Vec, - ) -> FlowResult<(String, mpsc::UnboundedReceiver)> { - let id = self.next_request_id(); - let (tx, rx) = mpsc::unbounded_channel(); - self.stream_events - .lock() - .map_err(|err| FlowError::Internal(format!("worker stream lock poisoned: {err}")))? - .insert(id.clone(), tx); - let payload = json!({ - "id": id, - "command": "stream_start", - "messages": messages, - }); - if let Err(err) = self.write_command(&payload) { - self.forget_stream(&id); - return Err(err); - } - Ok((id, rx)) - } - - fn send_stream_text(&self, id: &str, text: String) -> FlowResult<()> { - self.write_command(&json!({ - "id": id, - "command": "stream_text", - "text": text, - })) - } - - fn send_stream_end(&self, id: &str) -> FlowResult<()> { - self.write_command(&json!({ - "id": id, - "command": "stream_end", - })) - } - - fn forget_stream(&self, id: &str) { - let _ = self - .stream_events - .lock() - .map(|mut streams| streams.remove(id)); - } - - fn next_request_id(&self) -> String { - self.next_id.fetch_add(1, Ordering::Relaxed).to_string() - } - - fn write_command(&self, payload: &Json) -> FlowResult<()> { - let line = serde_json::to_string(payload).map_err(|err| { - FlowError::Internal(format!("failed to serialize worker command: {err}")) - })?; - let writer = self - .writer - .lock() - .map_err(|err| FlowError::Internal(format!("worker writer lock poisoned: {err}")))?; - writer - .as_ref() - .ok_or_else(|| FlowError::Internal("worker command writer is closed".to_string()))? - .send(line) - } - - fn shutdown(&self) { - if self.shutdown_started.swap(true, Ordering::AcqRel) { - return; - } - log::info!( - target: "nemo_relay.worker", - event = "worker_stopping", - plugin_id = "nemo_guardrails"; - "NeMo Guardrails local worker is stopping" - ); - let writer = self.writer.lock().ok().and_then(|mut writer| writer.take()); - let mut cleanup_succeeded = true; - if let Ok(mut child) = self.child.lock() { - if child.kill().is_err() { - cleanup_succeeded = false; - log::warn!( - target: "nemo_relay.worker", - event = "worker_cleanup_failed", - plugin_id = "nemo_guardrails", - operation = "kill"; - "NeMo Guardrails local worker cleanup failed" - ); - } - if child.wait().is_err() { - cleanup_succeeded = false; - log::warn!( - target: "nemo_relay.worker", - event = "worker_cleanup_failed", - plugin_id = "nemo_guardrails", - operation = "wait"; - "NeMo Guardrails local worker cleanup failed" - ); - } - } - if let Some(writer) = writer { - writer.join(); - } - if cleanup_succeeded { - log::info!( - target: "nemo_relay.worker", - event = "worker_stopped", - plugin_id = "nemo_guardrails"; - "NeMo Guardrails local worker stopped" - ); - } - } -} - -impl Drop for LocalGuardrailsWorker { - fn drop(&mut self) { - self.shutdown(); - } -} - -struct WorkerCommandWriter { - sender: std_mpsc::Sender, - error: Arc>>, - handle: Option>, -} - -impl WorkerCommandWriter { - fn spawn(mut stdin: ChildStdin) -> Self { - let (sender, receiver) = std_mpsc::channel::(); - let error = Arc::new(Mutex::new(None)); - let writer_error = Arc::clone(&error); - let handle = thread::spawn(move || { - for line in receiver { - if let Err(err) = writeln!(stdin, "{line}").and_then(|_| stdin.flush()) { - if let Ok(mut stored_error) = writer_error.lock() { - *stored_error = Some(err.to_string()); - } - return; - } - } - let _ = stdin.flush(); - }); - Self { - sender, - error, - handle: Some(handle), - } - } - - fn send(&self, line: String) -> FlowResult<()> { - if let Some(error) = self - .error - .lock() - .map_err(|err| { - FlowError::Internal(format!("worker writer error lock poisoned: {err}")) - })? - .clone() - { - return Err(FlowError::Internal(format!( - "failed to write worker command: {error}" - ))); - } - self.sender.send(line).map_err(|err| { - FlowError::Internal(format!("worker command writer channel closed: {err}")) - }) - } - - fn join(mut self) { - drop(self.sender); - if let Some(handle) = self.handle.take() { - let _ = handle.join(); - } - } -} - -#[derive(Debug, Clone, Deserialize)] -struct WorkerEnvelope { - id: String, - ok: bool, - #[serde(default)] - result: Option, - #[serde(default)] - error: Option, - #[serde(default)] - event: Option, - #[serde(default)] - message: Option, -} - -#[derive(Deserialize)] -struct WorkerCheckResult { - status: String, - #[serde(default)] - content: Option, - #[serde(default)] - rail: Option, -} - -fn python_executable(config: &NeMoGuardrailsConfig) -> String { - config - .local - .as_ref() - .and_then(|local| local.python_executable.as_deref()) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) - .or_else(|| env_executable(PYTHON_EXECUTABLE_ENV)) - .or_else(|| env_executable(PYO3_PYTHON_ENV)) - .or_else(|| env_executable(UV_PYTHON_ENV)) - .unwrap_or_else(|| DEFAULT_PYTHON_EXECUTABLE.to_string()) -} - -fn env_executable(name: &str) -> Option { - env::var(name) - .ok() - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) -} - -fn python_path(config: &NeMoGuardrailsConfig) -> Option { - config - .local - .as_ref() - .and_then(|local| local.python_path.as_deref()) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) -} - -fn worker_python_path(config: &NeMoGuardrailsConfig) -> Option { - let configured = python_path(config)?; - merge_python_path( - OsStr::new(&configured), - env::var_os("PYTHONPATH").as_deref(), - ) -} - -fn merge_python_path(configured: &OsStr, inherited: Option<&OsStr>) -> Option { - let mut paths = env::split_paths(configured).collect::>(); - if let Some(inherited) = inherited.filter(|value| !value.is_empty()) { - paths.extend(env::split_paths(inherited)); - } - env::join_paths(paths).ok() -} - -fn set_request_id(payload: &mut Json, id: &str) -> FlowResult<()> { - let object = payload.as_object_mut().ok_or_else(|| { - FlowError::Internal("worker command payload must be a JSON object".to_string()) - })?; - object.insert("id".to_string(), Json::String(id.to_string())); - Ok(()) -} - -fn dispatch_worker_envelope( - waiters: &Arc>>>, - stream_events: &Arc>>>, - envelope: WorkerEnvelope, -) { - if envelope.event.is_some() { - let sender = stream_events - .lock() - .ok() - .and_then(|streams| streams.get(&envelope.id).cloned()); - if let Some(sender) = sender { - let _ = sender.send(envelope); - } - return; - } - - let sender = waiters - .lock() - .ok() - .and_then(|mut waiters| waiters.remove(&envelope.id)); - if let Some(sender) = sender { - let _ = sender.send(envelope); - } -} - -fn notify_worker_closed( - waiters: &Arc>>>, - stream_events: &Arc>>>, - message: String, -) { - if let Ok(mut waiters) = waiters.lock() { - for (id, sender) in waiters.drain() { - let _ = sender.send(WorkerEnvelope { - id, - ok: false, - result: None, - error: Some(message.clone()), - event: None, - message: None, - }); - } - } - if let Ok(mut streams) = stream_events.lock() { - for (id, sender) in streams.drain() { - let _ = sender.send(WorkerEnvelope { - id, - ok: false, - result: None, - error: Some(message.clone()), - event: Some("error".to_string()), - message: None, - }); - } - } -} - -fn worker_result(envelope: WorkerEnvelope) -> FlowResult { - if envelope.ok { - Ok(envelope.result.unwrap_or(Json::Null)) - } else { - Err(FlowError::Internal(envelope.error.unwrap_or_else(|| { - "NeMo Guardrails local Python worker failed".to_string() - }))) - } -} - -fn parse_check_result(result: Json) -> FlowResult { - let result: WorkerCheckResult = serde_json::from_value(result).map_err(|err| { - FlowError::Internal(format!("worker returned invalid check result: {err}")) - })?; - match result.status.as_str() { - "blocked" => Ok(LocalCheckOutcome::Blocked { rail: result.rail }), - "modified" => Ok(LocalCheckOutcome::Modified { - content: result.content.unwrap_or_default(), - }), - "passed" => Ok(LocalCheckOutcome::Passed), - unexpected => Err(FlowError::Internal(format!( - "unexpected worker check status: {unexpected}" - ))), - } -} - -#[derive(Clone, Copy)] -enum LocalGuardrailsCodec { - OpenAIChat, - OpenAIResponses, - AnthropicMessages, - OCIGenAI, - GeminiGenerateContent, -} - -impl LocalGuardrailsCodec { - fn provider_surface(self) -> ProviderSurface { - match self { - Self::OpenAIChat => ProviderSurface::OpenAIChat, - Self::OpenAIResponses => ProviderSurface::OpenAIResponses, - Self::AnthropicMessages => ProviderSurface::AnthropicMessages, - Self::OCIGenAI => ProviderSurface::OCIGenAI, - Self::GeminiGenerateContent => ProviderSurface::GeminiGenerateContent, - } - } - - fn from_provider_surface(surface: ProviderSurface) -> Self { - match surface { - ProviderSurface::OpenAIChat => Self::OpenAIChat, - ProviderSurface::OpenAIResponses => Self::OpenAIResponses, - ProviderSurface::AnthropicMessages => Self::AnthropicMessages, - ProviderSurface::OCIGenAI => Self::OCIGenAI, - ProviderSurface::GeminiGenerateContent => Self::GeminiGenerateContent, - } - } - - fn decode(&self, request: &LlmRequest) -> FlowResult { - request_codec(self.provider_surface()).decode(request) - } - - fn encode( - &self, - annotated: &AnnotatedLlmRequest, - original: &LlmRequest, - ) -> FlowResult { - request_codec(self.provider_surface()).encode(annotated, original) - } - - fn decode_response( - &self, - response: &Json, - ) -> FlowResult { - response_codec(self.provider_surface()).decode_response(response) - } -} - -fn resolve_codec(config: &NeMoGuardrailsConfig) -> PluginResult> { - if !(config.input || config.output) { - return Ok(None); - } - - match config.codec.as_deref() { - Some(name) => match ProviderSurface::from_codec_name(name) { - Some(surface) => Ok(Some(LocalGuardrailsCodec::from_provider_surface(surface))), - None => Err(PluginError::InvalidConfig(format!( - "unsupported local NeMo Guardrails codec '{name}'" - ))), - }, - None => Err(PluginError::InvalidConfig( - "local NeMo Guardrails backend requires a supported codec".to_string(), - )), - } -} - -enum LocalCheckOutcome { - Passed, - Blocked { rail: Option }, - Modified { content: String }, -} - -#[derive(Clone, Copy)] -enum LocalRailKind { - Input, - Output, -} - -impl LocalRailKind { - fn as_str(self) -> &'static str { - match self { - Self::Input => "input", - Self::Output => "output", - } - } -} - -fn messages_from_annotated(annotated: &AnnotatedLlmRequest) -> FlowResult> { - match serde_json::to_value(&annotated.messages) - .map_err(|err| FlowError::Internal(format!("failed to serialize messages: {err}")))? - { - Json::Array(messages) => Ok(messages), - _ => Err(FlowError::Internal( - "serialized messages were not a JSON array".to_string(), - )), - } -} - -fn replace_last_role_content( - annotated: &mut AnnotatedLlmRequest, - role: &str, - content: String, -) -> FlowResult<()> { - for message in annotated.messages.iter_mut().rev() { - match (role, message) { - ( - "user", - Message::User { - content: target, .. - }, - ) => { - *target = MessageContent::Text(content); - return Ok(()); - } - ( - "assistant", - Message::Assistant { - content: target, .. - }, - ) => { - *target = Some(MessageContent::Text(content)); - return Ok(()); - } - _ => {} - } - } - - Err(local_violation(format!( - "NeMo Guardrails returned modified {role} content but no {role} message was present." - ))) -} - -fn tool_input_content(name: &str, args: &Json) -> FlowResult { - serde_json::to_string(&json!({ - "tool_name": name, - "arguments": args, - })) - .map_err(|err| FlowError::Internal(format!("failed to serialize tool input: {err}"))) -} - -fn tool_output_content(name: &str, args: &Json, result: &Json) -> FlowResult { - serde_json::to_string(&json!({ - "tool_name": name, - "arguments": args, - "result": result, - })) - .map_err(|err| FlowError::Internal(format!("failed to serialize tool output: {err}"))) -} - -fn modified_tool_payload(content: &str, field: &str) -> FlowResult { - let value: Json = serde_json::from_str(content).map_err(|_| { - local_violation(format!( - "NeMo Guardrails returned modified tool {field} content that is not valid JSON." - )) - })?; - - let Json::Object(object) = value else { - return Err(local_violation(format!( - "NeMo Guardrails returned modified tool {field} content without a '{field}' field." - ))); - }; - object.get(field).cloned().ok_or_else(|| { - local_violation(format!( - "NeMo Guardrails returned modified tool {field} content without a '{field}' field." - )) - }) -} - -fn blocked_error(rail_type: &str, rail: Option<&str>) -> FlowError { - let detail = rail - .filter(|rail| !rail.is_empty()) - .map(|rail| format!(" by rail '{rail}'")) - .unwrap_or_default(); - let subject = if matches!(rail_type, "input" | "output") { - "LLM call" - } else { - "tool call" - }; - local_violation(format!( - "NeMo Guardrails {rail_type} rail blocked the {subject}{detail}." - )) -} - -fn local_violation(message: impl Into) -> FlowError { - FlowError::Internal(message.into()) -} - -struct GuardedProviderStream { - receiver: ReceiverStream>, - cancel: watch::Sender, - closed: watch::Receiver>>, -} - -impl tokio_stream::Stream for GuardedProviderStream { - type Item = FlowResult; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.receiver).poll_next(cx) - } -} - -impl Drop for GuardedProviderStream { - fn drop(&mut self) { - self.cancel.send_replace(true); - } -} - -impl LlmStreamInner for GuardedProviderStream { - fn close(self: Pin<&mut Self>) -> Pin> + Send + '_>> { - let this = self.get_mut(); - this.cancel.send_replace(true); - this.receiver.close(); - while this.receiver.as_mut().try_recv().is_ok() {} - let mut closed = this.closed.clone(); - Box::pin(async move { - while closed.borrow().is_none() { - closed.changed().await.map_err(|_| { - FlowError::Internal("guarded stream cleanup task ended early".into()) - })?; - } - closed.borrow().clone().expect("close state checked above") - }) - } -} - -#[allow( - clippy::too_many_arguments, - reason = "stream cancellation, monitoring, delivery, and cleanup must remain ordered in one coordinator" -)] -async fn forward_guarded_provider_stream( - mut provider_stream: LlmJsonStream, - codec: LocalGuardrailsCodec, - text_tx: mpsc::Sender>, - chunk_tx: mpsc::Sender>, - monitor: JoinHandle>, - blocked: Arc>>, - mut cancel: watch::Receiver, - closed: watch::Sender>>, -) { - let mut monitor = Some(monitor); - loop { - if *cancel.borrow() { - break; - } - let item = tokio::select! { - _ = cancel.changed() => break, - item = provider_stream.next() => item, - }; - let Some(item) = item else { - break; - }; - let Some(chunk) = - receive_guarded_provider_chunk(item, &text_tx, &chunk_tx, &mut monitor).await - else { - break; - }; - - if stop_blocked_provider_stream(&text_tx, &chunk_tx, &blocked, &mut monitor).await { - break; - } - if !forward_guarded_stream_text(codec, &chunk, &text_tx, &chunk_tx, &blocked, &mut monitor) - .await - { - break; - } - - if !send_guarded_provider_chunk(chunk, &text_tx, &chunk_tx, &mut monitor, &mut cancel).await - { - break; - } - } - finish_guarded_provider_stream( - &mut provider_stream, - &text_tx, - &chunk_tx, - &blocked, - &mut monitor, - &cancel, - &closed, - ) - .await; -} - -async fn receive_guarded_provider_chunk( - item: FlowResult, - text_tx: &mpsc::Sender>, - chunk_tx: &mpsc::Sender>, - monitor: &mut Option>>, -) -> Option { - match item { - Ok(chunk) => Some(chunk), - Err(err) => { - let _ = chunk_tx.send(Err(err)).await; - let _ = text_tx.send(None).await; - let _ = monitor.take().expect("monitor available").await; - None - } - } -} - -async fn stop_blocked_provider_stream( - text_tx: &mpsc::Sender>, - chunk_tx: &mpsc::Sender>, - blocked: &Arc>>, - monitor: &mut Option>>, -) -> bool { - let Some(message) = blocked_message(blocked) else { - return false; - }; - let _ = chunk_tx.send(Err(streaming_output_blocked(message))).await; - let _ = text_tx.send(None).await; - let _ = monitor.take().expect("monitor available").await; - true -} - -async fn forward_guarded_stream_text( - codec: LocalGuardrailsCodec, - chunk: &Json, - text_tx: &mpsc::Sender>, - chunk_tx: &mpsc::Sender>, - blocked: &Arc>>, - monitor: &mut Option>>, -) -> bool { - let Some(text) = extract_stream_text(codec, chunk) else { - return true; - }; - if text_tx.send(Some(text)).await.is_ok() { - return true; - } - send_stream_monitor_error( - monitor.take().expect("monitor available"), - chunk_tx, - blocked, - ) - .await; - false -} - -async fn send_guarded_provider_chunk( - chunk: Json, - text_tx: &mpsc::Sender>, - chunk_tx: &mpsc::Sender>, - monitor: &mut Option>>, - cancel: &mut watch::Receiver, -) -> bool { - let sent = tokio::select! { - _ = cancel.changed() => return false, - sent = chunk_tx.send(Ok(chunk)) => sent, - }; - if sent.is_ok() { - return true; - } - let _ = text_tx.send(None).await; - let _ = monitor.take().expect("monitor available").await; - false -} - -#[allow( - clippy::too_many_arguments, - reason = "stream cleanup needs all channels and lifecycle handles" -)] -async fn finish_guarded_provider_stream( - provider_stream: &mut LlmJsonStream, - text_tx: &mpsc::Sender>, - chunk_tx: &mpsc::Sender>, - blocked: &Arc>>, - monitor: &mut Option>>, - cancel: &watch::Receiver, - closed: &watch::Sender>>, -) { - let _ = text_tx.send(None).await; - if *cancel.borrow() { - if let Some(monitor) = monitor.take() { - monitor.abort(); - } - } else if let Some(monitor) = monitor.take() { - let _ = send_stream_monitor_error(monitor, chunk_tx, blocked).await; - } - closed.send_replace(Some(provider_stream.close().await)); -} - -async fn send_stream_monitor_error( - monitor: JoinHandle>, - chunk_tx: &mpsc::Sender>, - blocked: &Arc>>, -) -> bool { - match monitor.await { - Ok(Ok(())) => {} - Ok(Err(err)) => { - let _ = chunk_tx.send(Err(err)).await; - return true; - } - Err(err) => { - let _ = chunk_tx - .send(Err(FlowError::Internal(format!( - "nemo_guardrails stream monitor task failed: {err}" - )))) - .await; - return true; - } - } - - if let Some(message) = blocked_message(blocked) { - let _ = chunk_tx.send(Err(streaming_output_blocked(message))).await; - return true; - } - - false -} - -fn blocked_message(blocked: &Arc>>) -> Option { - blocked.lock().ok().and_then(|guard| guard.clone()) -} - -fn streaming_output_blocked(message: String) -> FlowError { - local_violation(format!( - "NeMo Guardrails output rail blocked the LLM call: {message}" - )) -} - -fn extract_stream_text(codec: LocalGuardrailsCodec, chunk: &Json) -> Option { - let chunk = chunk.as_object()?; - match codec { - LocalGuardrailsCodec::OpenAIChat => extract_openai_chat_stream_text(chunk), - LocalGuardrailsCodec::OpenAIResponses => extract_openai_response_stream_text(chunk), - LocalGuardrailsCodec::AnthropicMessages => extract_anthropic_stream_text(chunk), - LocalGuardrailsCodec::OCIGenAI => extract_oci_genai_stream_text(chunk), - LocalGuardrailsCodec::GeminiGenerateContent => extract_gemini_stream_text(chunk), - } -} - -/// Collect the concatenated TEXT-part text from OCI GENERIC stream deltas or -/// the bare `text` fragment of COHERE deltas. Events may arrive wrapped in a -/// `chatResponse` envelope, and GENERIC deltas are either a bare choice -/// (`message` at the top level) or carry a `choices` array of deltas, -/// mirroring the stream shapes the OCI streaming codec accepts. -/// -/// The live service's terminal COHERE event (the one carrying `finishReason`) -/// repeats the complete response text already delivered by earlier deltas. -/// Forwarding it would double the text the output rails evaluate, so it is -/// suppressed here, mirroring the deduplication in the OCI streaming codec. -fn extract_oci_genai_stream_text(chunk: &serde_json::Map) -> Option { - let chunk = chunk - .get("chatResponse") - .and_then(Json::as_object) - .unwrap_or(chunk); - if let Some(text) = chunk.get("text").and_then(Json::as_str) { - if chunk.get("finishReason").and_then(Json::as_str).is_some() { - return None; - } - return (!text.is_empty()).then(|| text.to_string()); - } - fn collect_generic_text(message: &Json, collected: &mut String) { - let Some(parts) = message.get("content").and_then(Json::as_array) else { - return; - }; - for part in parts { - if part.get("type").and_then(Json::as_str) == Some("TEXT") - && let Some(text) = part.get("text").and_then(Json::as_str) - { - collected.push_str(text); - } - } - } - let mut collected = String::new(); - match chunk.get("choices").and_then(Json::as_array) { - Some(choices) => { - for choice in choices { - if let Some(message) = choice.get("message") { - collect_generic_text(message, &mut collected); - } - } - } - None => { - if let Some(message) = chunk.get("message") { - collect_generic_text(message, &mut collected); - } - } - } - (!collected.is_empty()).then_some(collected) -} - -fn extract_openai_chat_stream_text(chunk: &serde_json::Map) -> Option { - let choices = chunk.get("choices")?.as_array()?; - let parts = choices - .iter() - .filter_map(|choice| { - choice - .get("delta") - .and_then(Json::as_object) - .and_then(|delta| delta.get("content")) - .and_then(Json::as_str) - .filter(|content| !content.is_empty()) - }) - .collect::>(); - (!parts.is_empty()).then(|| parts.join("")) -} - -fn extract_openai_response_stream_text(chunk: &serde_json::Map) -> Option { - (chunk.get("type").and_then(Json::as_str) == Some("response.output_text.delta")) - .then(|| chunk.get("delta").and_then(Json::as_str)) - .flatten() - .filter(|delta| !delta.is_empty()) - .map(str::to_string) -} - -fn extract_anthropic_stream_text(chunk: &serde_json::Map) -> Option { - if chunk.get("type").and_then(Json::as_str) != Some("content_block_delta") { - return None; - } - let delta = chunk.get("delta")?.as_object()?; - (delta.get("type").and_then(Json::as_str) == Some("text_delta")) - .then(|| delta.get("text").and_then(Json::as_str)) - .flatten() - .filter(|text| !text.is_empty()) - .map(str::to_string) -} - -fn extract_gemini_stream_text(chunk: &serde_json::Map) -> Option { - let parts = chunk - .get("candidates")? - .as_array()? - .first()? - .get("content")? - .get("parts")? - .as_array()?; - let texts = parts - .iter() - .filter(|part| part.get("thought").and_then(Json::as_bool) != Some(true)) - .filter_map(|part| part.get("text").and_then(Json::as_str)) - .filter(|text| !text.is_empty()) - .collect::>(); - (!texts.is_empty()).then(|| texts.join("")) -} - -async fn monitor_guardrails_stream( - worker: Arc, - stream_id: String, - mut text_rx: mpsc::Receiver>, - mut event_rx: mpsc::UnboundedReceiver, - blocked: Arc>>, -) -> FlowResult<()> { - let mut input_closed = false; - loop { - tokio::select! { - maybe_text = text_rx.recv(), if !input_closed => { - match maybe_text { - Some(Some(text)) => worker.send_stream_text(&stream_id, text)?, - Some(None) | None => { - worker.send_stream_end(&stream_id)?; - input_closed = true; - } - } - } - maybe_event = event_rx.recv() => { - let Some(event) = maybe_event else { - worker.forget_stream(&stream_id); - return Err(FlowError::Internal( - "NeMo Guardrails local Python worker stream closed unexpectedly".to_string(), - )); - }; - if !event.ok { - worker.forget_stream(&stream_id); - return Err(FlowError::Internal(event.error.unwrap_or_else(|| { - "NeMo Guardrails local Python worker stream failed".to_string() - }))); - } - match event.event.as_deref() { - Some("blocked") => { - if let Some(message) = event.message { - let mut guard = blocked.lock().map_err(|err| { - FlowError::Internal(format!("stream block state lock poisoned: {err}")) - })?; - *guard = Some(message); - } - worker.forget_stream(&stream_id); - return Ok(()); - } - Some("done") => { - worker.forget_stream(&stream_id); - return Ok(()); - } - Some(other) => { - worker.forget_stream(&stream_id); - return Err(FlowError::Internal(format!( - "NeMo Guardrails local Python worker returned unknown stream event '{other}'" - ))); - } - None => {} - } - } - } - } -} - -#[cfg(test)] -#[path = "../../../tests/unit/plugins/nemo_guardrails/local_python_tests.rs"] -mod tests; diff --git a/crates/core/src/plugins/nemo_guardrails/remote.rs b/crates/core/src/plugins/nemo_guardrails/remote.rs deleted file mode 100644 index d349cd1514..0000000000 --- a/crates/core/src/plugins/nemo_guardrails/remote.rs +++ /dev/null @@ -1,1203 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use std::sync::Arc; -use std::sync::atomic::{AtomicU8, Ordering}; -use std::time::Duration; - -use serde_json::{Map, Value as Json, json}; -use tokio::sync::mpsc; -use tokio_stream::wrappers::ReceiverStream; - -use crate::api::llm::LlmRequest; -use crate::api::runtime::{LlmExecutionFn, LlmJsonStream, LlmStreamExecutionFn, ToolExecutionFn}; -use crate::api::scope::{EmitMarkEventParams, ScopeHandle, event, get_handle}; -use crate::codec::openai_chat::OpenAIChatCodec; -use crate::codec::streaming::SseEventDecoder; -use crate::codec::traits::LlmCodec; -use crate::error::FlowError; -use crate::plugin::{PluginError, PluginRegistrationContext, Result as PluginResult}; -use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; - -use super::{NeMoGuardrailsConfig, RequestDefaultsConfig, RequestRailsConfig}; - -#[derive(Clone)] -struct RemoteBackendRuntime { - endpoint: String, - client: reqwest::Client, - config_id: Option, - config_ids: Vec, - llm_guardrails: Option>, - tool_input_guardrails: Map, - tool_output_guardrails: Map, - access_state: Arc, -} - -#[derive(Clone, Copy)] -enum RemoteCheckKind { - Input, - Output, -} - -impl RemoteBackendRuntime { - fn new(config: &NeMoGuardrailsConfig) -> PluginResult { - let remote = config.remote.as_ref().ok_or_else(|| { - PluginError::InvalidConfig( - "remote config is required when mode is 'remote'".to_string(), - ) - })?; - let endpoint = remote.endpoint.clone().ok_or_else(|| { - PluginError::InvalidConfig("remote.endpoint is required in remote mode".to_string()) - })?; - let mut default_headers = HeaderMap::new(); - for (name, value) in &remote.headers { - let header_name = HeaderName::from_bytes(name.as_bytes()).map_err(|err| { - PluginError::InvalidConfig(format!( - "remote.headers contains invalid header name '{name}': {err}" - )) - })?; - let header_value = HeaderValue::from_str(value).map_err(|err| { - PluginError::InvalidConfig(format!( - "remote.headers[{name}] has an invalid value: {err}" - )) - })?; - default_headers.insert(header_name, header_value); - } - - let client = reqwest::Client::builder() - .default_headers(default_headers) - .timeout(Duration::from_millis(remote.timeout_millis)) - .build() - .map_err(|err| { - PluginError::RegistrationFailed(format!( - "failed to construct NeMo Guardrails remote client: {err}" - )) - })?; - - let request_defaults = config.request_defaults.as_ref(); - - log::info!( - target: "nemo_relay.plugin", - event = "plugin_resource_access_pending", - plugin_kind = "nemo_guardrails", - resource_kind = "http_endpoint", - permission = "invoke"; - "Plugin resource access will be validated on first use" - ); - - Ok(Self { - endpoint: endpoint.trim_end_matches('/').to_string(), - client, - config_id: remote.config_id.clone(), - config_ids: remote.config_ids.clone(), - llm_guardrails: build_llm_guardrails_config( - &remote.config_id, - &remote.config_ids, - request_defaults, - config.input, - config.output, - ), - tool_input_guardrails: build_tool_check_guardrails_config( - RemoteCheckKind::Input, - &remote.config_id, - &remote.config_ids, - request_defaults, - ), - tool_output_guardrails: build_tool_check_guardrails_config( - RemoteCheckKind::Output, - &remote.config_id, - &remote.config_ids, - request_defaults, - ), - access_state: Arc::new(AtomicU8::new(0)), - }) - } - - async fn execute(&self, request: LlmRequest, stream: bool) -> crate::error::Result { - let parent = get_handle().ok(); - self.emit_remote_start(&parent, stream); - let body = self.build_request_body_with_marks(&parent, &request, stream)?; - let response = self - .send_remote_request_with_marks(&parent, stream, body) - .await?; - let status = response.status(); - let response_json = self - .read_json_response_with_marks(&parent, stream, response) - .await?; - self.emit_mark( - "nemo_guardrails.remote.end", - &parent, - remote_mark_data( - stream, - &self.config_id, - &self.config_ids, - Some(status.as_u16()), - None, - ), - ); - Ok(response_json) - } - - async fn execute_stream(&self, request: LlmRequest) -> crate::error::Result { - let parent = get_handle().ok(); - self.emit_remote_start(&parent, true); - let body = self.build_request_body_with_marks(&parent, &request, true)?; - let response = self - .send_remote_request_with_marks(&parent, true, body) - .await?; - let status = response.status(); - if !status.is_success() { - let payload = response.text().await.map_err(|err| { - self.emit_mark( - "nemo_guardrails.remote.error", - &parent, - remote_mark_data( - true, - &self.config_id, - &self.config_ids, - Some(status.as_u16()), - Some(format!("failed to read remote stream error body: {err}")), - ), - ); - FlowError::Internal(format!( - "nemo_guardrails failed to read remote stream error body: {err}" - )) - })?; - self.emit_mark( - "nemo_guardrails.remote.error", - &parent, - remote_mark_data( - true, - &self.config_id, - &self.config_ids, - Some(status.as_u16()), - Some(redact_remote_error_payload(status.as_u16(), &payload)), - ), - ); - return Err(FlowError::Internal(format!( - "nemo_guardrails remote stream request failed with status {status}: {payload}" - ))); - } - - let (tx, rx) = mpsc::channel(16); - self.spawn_stream_decoder(response, status, parent.clone(), tx); - - Ok(LlmJsonStream::new(ReceiverStream::new(rx))) - } - - fn emit_remote_start(&self, parent: &Option, stream: bool) { - self.emit_mark( - "nemo_guardrails.remote.start", - parent, - remote_mark_data(stream, &self.config_id, &self.config_ids, None, None), - ); - } - - fn emit_remote_error( - &self, - parent: &Option, - stream: bool, - status: Option, - error: impl Into, - ) { - self.emit_mark( - "nemo_guardrails.remote.error", - parent, - remote_mark_data( - stream, - &self.config_id, - &self.config_ids, - status, - Some(error.into()), - ), - ); - } - - fn build_request_body_with_marks( - &self, - parent: &Option, - request: &LlmRequest, - stream: bool, - ) -> crate::error::Result { - self.build_request_body(request, stream).inspect_err(|err| { - self.emit_remote_error(parent, stream, None, err.to_string()); - }) - } - - async fn send_remote_request_with_marks( - &self, - parent: &Option, - stream: bool, - body: Json, - ) -> crate::error::Result { - let serialized = self.serialize_request_body_with_marks(parent, stream, body)?; - match self - .client - .post(self.chat_completions_url()) - .header(reqwest::header::CONTENT_TYPE, "application/json") - .body(serialized) - .send() - .await - { - Ok(response) => { - self.record_access_status(response.status()); - Ok(response) - } - Err(err) => { - self.record_access_failure("connection_failed", None); - let message = if stream { - format!("remote stream request failed: {err}") - } else { - format!("remote request failed: {err}") - }; - self.emit_remote_error(parent, stream, None, message.clone()); - Err(FlowError::Internal(format!("nemo_guardrails {message}"))) - } - } - } - - fn record_access_status(&self, status: reqwest::StatusCode) { - if status.is_success() { - if self.access_state.swap(2, Ordering::AcqRel) != 2 { - log::info!( - target: "nemo_relay.plugin", - event = "plugin_resource_access_validated", - plugin_kind = "nemo_guardrails", - resource_kind = "http_endpoint", - permission = "invoke", - status_code = status.as_u16(); - "Plugin resource access validated" - ); - } - return; - } - let reason = if matches!(status.as_u16(), 401 | 403) { - "permission_denied" - } else { - "unsuccessful_status" - }; - self.record_access_failure(reason, Some(status.as_u16())); - } - - fn record_access_failure(&self, reason: &'static str, status_code: Option) { - if self.access_state.swap(1, Ordering::AcqRel) == 1 { - return; - } - match status_code { - Some(status_code) => log::warn!( - target: "nemo_relay.plugin", - event = "plugin_resource_access_failed", - plugin_kind = "nemo_guardrails", - resource_kind = "http_endpoint", - permission = "invoke", - reason = reason, - status_code = status_code; - "Plugin resource access validation failed" - ), - None => log::warn!( - target: "nemo_relay.plugin", - event = "plugin_resource_access_failed", - plugin_kind = "nemo_guardrails", - resource_kind = "http_endpoint", - permission = "invoke", - reason = reason; - "Plugin resource access validation failed" - ), - } - } - - fn serialize_request_body_with_marks( - &self, - parent: &Option, - stream: bool, - body: Json, - ) -> crate::error::Result> { - serde_json::to_vec(&body).map_err(|err| { - let context = if stream { - "remote stream request body" - } else { - "remote request body" - }; - let message = format!("failed to serialize {context}: {err}"); - self.emit_remote_error(parent, stream, None, message.clone()); - FlowError::Internal(format!("nemo_guardrails {message}")) - }) - } - - async fn read_json_response_with_marks( - &self, - parent: &Option, - stream: bool, - response: reqwest::Response, - ) -> crate::error::Result { - let status = response.status(); - let payload = self - .read_response_text_with_marks(parent, stream, response, status) - .await?; - if !status.is_success() { - self.emit_remote_error( - parent, - stream, - Some(status.as_u16()), - redact_remote_error_payload(status.as_u16(), &payload), - ); - return Err(FlowError::Internal(format!( - "nemo_guardrails remote request failed with status {status}: {payload}" - ))); - } - - serde_json::from_str(&payload).map_err(|err| { - let message = format!("failed to parse remote response JSON: {err}"); - self.emit_remote_error(parent, stream, Some(status.as_u16()), message.clone()); - FlowError::Internal(format!("nemo_guardrails {message}")) - }) - } - - async fn read_response_text_with_marks( - &self, - parent: &Option, - stream: bool, - response: reqwest::Response, - status: reqwest::StatusCode, - ) -> crate::error::Result { - response.text().await.map_err(|err| { - let context = if stream { - "remote stream error body" - } else { - "remote response body" - }; - let message = format!("failed to read {context}: {err}"); - self.emit_remote_error(parent, stream, Some(status.as_u16()), message.clone()); - FlowError::Internal(format!("nemo_guardrails {message}")) - }) - } - - fn spawn_stream_decoder( - &self, - mut response: reqwest::Response, - status: reqwest::StatusCode, - parent: Option, - tx: mpsc::Sender>, - ) { - let config_id = self.config_id.clone(); - let config_ids = self.config_ids.clone(); - tokio::spawn(async move { - let mut decoder = SseEventDecoder::new(); - loop { - let bytes = match response.chunk().await { - Ok(Some(bytes)) => bytes, - Ok(None) => break, - Err(err) => { - emit_stream_decode_error( - &parent, - &config_id, - &config_ids, - status, - format!("failed to read remote stream chunk: {err}"), - ); - let _ = tx - .send(Err(FlowError::Internal(format!( - "nemo_guardrails failed to read remote stream chunk: {err}" - )))) - .await; - return; - } - }; - let events = match decoder.push_bytes(&bytes) { - Ok(events) => events, - Err(err) => { - emit_stream_decode_error( - &parent, - &config_id, - &config_ids, - status, - err.to_string(), - ); - let _ = tx.send(Err(err)).await; - return; - } - }; - for event in events { - if tx.send(Ok(event.data)).await.is_err() { - return; - } - } - } - - match decoder.finish() { - Ok(Some(event)) => { - let _ = tx.send(Ok(event.data)).await; - } - Ok(None) => {} - Err(err) => { - emit_stream_decode_error( - &parent, - &config_id, - &config_ids, - status, - err.to_string(), - ); - let _ = tx.send(Err(err)).await; - return; - } - } - - emit_remote_mark( - "nemo_guardrails.remote.end", - &parent, - remote_mark_data(true, &config_id, &config_ids, Some(status.as_u16()), None), - ); - }); - } - - async fn check_tool_input(&self, tool_name: &str, args: &Json) -> crate::error::Result { - let messages = tool_input_messages(tool_name, args); - let response = self - .execute_remote_check(messages, RemoteCheckKind::Input, tool_name) - .await?; - if let Some(blocking_rail) = blocking_rail_name(&response) { - return Err(FlowError::GuardrailRejected(format!( - "nemo_guardrails tool_input rail blocked tool call by rail '{blocking_rail}'" - ))); - } - - if let Some(modified_args) = modified_tool_arguments(&response, tool_name)? { - return Ok(modified_args); - } - Ok(args.clone()) - } - - async fn check_tool_output( - &self, - tool_name: &str, - args: &Json, - result: &Json, - ) -> crate::error::Result { - let messages = tool_output_messages(tool_name, args, result); - let response = self - .execute_remote_check(messages, RemoteCheckKind::Output, tool_name) - .await?; - if let Some(blocking_rail) = blocking_rail_name(&response) { - return Err(FlowError::GuardrailRejected(format!( - "nemo_guardrails tool_output rail blocked tool call by rail '{blocking_rail}'" - ))); - } - - if let Some(modified_result) = modified_tool_result(&response, tool_name)? { - return Ok(modified_result); - } - Ok(result.clone()) - } - - fn build_request_body(&self, request: &LlmRequest, stream: bool) -> crate::error::Result { - let annotated = OpenAIChatCodec.decode(request)?; - if annotated.tools.is_some() || annotated.tool_choice.is_some() { - return Err(FlowError::Internal( - "nemo_guardrails remote backend does not support OpenAI tool definitions or tool_choice yet" - .to_string(), - )); - } - - let mut body = request.content.as_object().cloned().ok_or_else(|| { - FlowError::Internal("LLM request content is not a JSON object".to_string()) - })?; - body.insert("stream".to_string(), Json::Bool(stream)); - if let Some(guardrails) = &self.llm_guardrails { - body.insert("guardrails".to_string(), Json::Object(guardrails.clone())); - } - Ok(Json::Object(body)) - } - - fn chat_completions_url(&self) -> String { - format!("{}/v1/chat/completions", self.endpoint) - } - - fn emit_mark(&self, name: &str, parent: &Option, data: Json) { - emit_remote_mark(name, parent, data); - } - - async fn execute_remote_check( - &self, - messages: Vec, - kind: RemoteCheckKind, - tool_name: &str, - ) -> crate::error::Result { - let parent = get_handle().ok(); - self.emit_mark( - "nemo_guardrails.remote.start", - &parent, - tool_remote_mark_data( - kind, - tool_name, - &self.config_id, - &self.config_ids, - None, - None, - ), - ); - let mut body = Map::new(); - body.insert("model".to_string(), Json::String(String::new())); - body.insert("messages".to_string(), Json::Array(messages)); - body.insert("stream".to_string(), Json::Bool(false)); - body.insert( - "guardrails".to_string(), - Json::Object(match kind { - RemoteCheckKind::Input => self.tool_input_guardrails.clone(), - RemoteCheckKind::Output => self.tool_output_guardrails.clone(), - }), - ); - let serialized = serde_json::to_vec(&Json::Object(body)).map_err(|err| { - let message = format!("nemo_guardrails failed to serialize remote request body: {err}"); - self.emit_mark( - "nemo_guardrails.remote.error", - &parent, - tool_remote_mark_data( - kind, - tool_name, - &self.config_id, - &self.config_ids, - None, - Some(message.clone()), - ), - ); - FlowError::Internal(message) - })?; - let response = match self - .client - .post(self.chat_completions_url()) - .header(reqwest::header::CONTENT_TYPE, "application/json") - .body(serialized) - .send() - .await - { - Ok(response) => { - self.record_access_status(response.status()); - response - } - Err(err) => { - self.record_access_failure("connection_failed", None); - let message = format!("nemo_guardrails remote request failed: {err}"); - self.emit_mark( - "nemo_guardrails.remote.error", - &parent, - tool_remote_mark_data( - kind, - tool_name, - &self.config_id, - &self.config_ids, - None, - Some(message.clone()), - ), - ); - return Err(FlowError::Internal(message)); - } - }; - let status = response.status(); - let payload = response.text().await.map_err(|err| { - let message = format!("nemo_guardrails failed to read remote response body: {err}"); - self.emit_mark( - "nemo_guardrails.remote.error", - &parent, - tool_remote_mark_data( - kind, - tool_name, - &self.config_id, - &self.config_ids, - Some(status.as_u16()), - Some(message.clone()), - ), - ); - FlowError::Internal(message) - })?; - if !status.is_success() { - self.emit_mark( - "nemo_guardrails.remote.error", - &parent, - tool_remote_mark_data( - kind, - tool_name, - &self.config_id, - &self.config_ids, - Some(status.as_u16()), - Some(redact_remote_error_payload(status.as_u16(), &payload)), - ), - ); - return Err(FlowError::Internal(format!( - "nemo_guardrails remote request failed with status {status}: {payload}" - ))); - } - let response_json = serde_json::from_str(&payload).map_err(|err| { - let message = format!("nemo_guardrails failed to parse remote response JSON: {err}"); - self.emit_mark( - "nemo_guardrails.remote.error", - &parent, - tool_remote_mark_data( - kind, - tool_name, - &self.config_id, - &self.config_ids, - Some(status.as_u16()), - Some(message.clone()), - ), - ); - FlowError::Internal(message) - })?; - self.emit_mark( - "nemo_guardrails.remote.end", - &parent, - tool_remote_mark_data( - kind, - tool_name, - &self.config_id, - &self.config_ids, - Some(status.as_u16()), - None, - ), - ); - Ok(response_json) - } -} - -fn emit_remote_mark(name: &str, parent: &Option, data: Json) { - let _ = event( - EmitMarkEventParams::builder() - .name(name) - .parent_opt(parent.as_ref()) - .data(data) - .build(), - ); -} - -fn emit_stream_decode_error( - parent: &Option, - config_id: &Option, - config_ids: &[String], - status: reqwest::StatusCode, - error: String, -) { - emit_remote_mark( - "nemo_guardrails.remote.error", - parent, - remote_mark_data( - true, - config_id, - config_ids, - Some(status.as_u16()), - Some(error), - ), - ); -} - -fn tool_call_id(tool_name: &str) -> String { - format!("nemo_guardrails_{tool_name}_call") -} - -fn redact_remote_error_payload(status: u16, payload: &str) -> String { - format!( - "remote request failed with status {status}; error body omitted from marks ({} bytes)", - payload.len() - ) -} - -fn tool_arguments_string(args: &Json) -> String { - serde_json::to_string(args).expect("tool arguments should serialize to JSON") -} - -fn tool_result_string(result: &Json) -> String { - serde_json::to_string(result).expect("tool result should serialize to JSON") -} - -fn tool_user_message(tool_name: &str) -> Json { - json!({ - "role": "user", - "content": format!("Run the tool '{tool_name}' and validate the result."), - }) -} - -fn tool_input_messages(tool_name: &str, args: &Json) -> Vec { - vec![ - tool_user_message(tool_name), - json!({ - "role": "assistant", - "content": "", - "tool_calls": [{ - "id": tool_call_id(tool_name), - "type": "function", - "function": { - "name": tool_name, - "arguments": tool_arguments_string(args), - } - }] - }), - ] -} - -fn tool_output_messages(tool_name: &str, args: &Json, result: &Json) -> Vec { - let call_id = tool_call_id(tool_name); - vec![ - tool_user_message(tool_name), - json!({ - "role": "assistant", - "content": "", - "tool_calls": [{ - "id": call_id, - "type": "function", - "function": { - "name": tool_name, - "arguments": tool_arguments_string(args), - } - }] - }), - json!({ - "role": "tool", - "name": tool_name, - "tool_call_id": call_id, - "content": tool_result_string(result), - }), - ] -} - -fn first_choice_message(response: &Json) -> crate::error::Result<&Map> { - response - .get("choices") - .and_then(Json::as_array) - .and_then(|choices| choices.first()) - .and_then(|choice| choice.get("message")) - .and_then(Json::as_object) - .ok_or_else(|| { - FlowError::Internal( - "nemo_guardrails remote response did not contain choices[0].message".to_string(), - ) - }) -} - -fn first_tool_call_message(message: &Map) -> Option<&Map> { - message - .get("tool_calls") - .and_then(Json::as_array) - .and_then(|tool_calls| tool_calls.first()) - .and_then(Json::as_object) -} - -fn modified_tool_arguments( - response: &Json, - expected_tool_name: &str, -) -> crate::error::Result> { - let message = first_choice_message(response)?; - if let Some(tool_call) = first_tool_call_message(message) { - let function = tool_call - .get("function") - .and_then(Json::as_object) - .ok_or_else(|| { - FlowError::Internal( - "nemo_guardrails returned modified tool arguments without a function payload" - .to_string(), - ) - })?; - let tool_name = function.get("name").and_then(Json::as_str).ok_or_else(|| { - FlowError::Internal( - "nemo_guardrails returned modified tool arguments without a function name" - .to_string(), - ) - })?; - if tool_name != expected_tool_name { - return Err(FlowError::Internal(format!( - "nemo_guardrails returned modified tool arguments for unexpected tool '{tool_name}'" - ))); - } - let arguments = function - .get("arguments") - .and_then(Json::as_str) - .ok_or_else(|| { - FlowError::Internal( - "nemo_guardrails returned modified tool arguments without function.arguments" - .to_string(), - ) - })?; - let parsed = serde_json::from_str(arguments).map_err(|err| { - FlowError::Internal(format!( - "nemo_guardrails returned modified tool arguments that are not valid JSON: {err}" - )) - })?; - return Ok(Some(parsed)); - } - - let content = message - .get("content") - .and_then(Json::as_str) - .filter(|content| !content.is_empty()); - legacy_modified_tool_payload(content, expected_tool_name, "arguments") -} - -fn modified_tool_result( - response: &Json, - expected_tool_name: &str, -) -> crate::error::Result> { - let message = first_choice_message(response)?; - if message.get("role").and_then(Json::as_str) == Some("tool") { - if let Some(tool_name) = message.get("name").and_then(Json::as_str) - && tool_name != expected_tool_name - { - return Err(FlowError::Internal(format!( - "nemo_guardrails returned modified tool result for unexpected tool '{tool_name}'" - ))); - } - let content = message - .get("content") - .and_then(Json::as_str) - .ok_or_else(|| { - FlowError::Internal( - "nemo_guardrails returned modified tool result without message.content" - .to_string(), - ) - })?; - let parsed = serde_json::from_str(content).map_err(|err| { - FlowError::Internal(format!( - "nemo_guardrails returned modified tool result that is not valid JSON: {err}" - )) - })?; - return Ok(Some(parsed)); - } - - let content = message - .get("content") - .and_then(Json::as_str) - .filter(|content| !content.is_empty()); - legacy_modified_tool_payload(content, expected_tool_name, "result") -} - -fn legacy_modified_tool_payload( - content: Option<&str>, - expected_tool_name: &str, - field: &str, -) -> crate::error::Result> { - let Some(content) = content else { - return Ok(None); - }; - let Ok(value) = serde_json::from_str(content) else { - return Ok(None); - }; - let Json::Object(object) = value else { - return Ok(None); - }; - if let Some(tool_name) = object.get("tool_name").and_then(Json::as_str) - && tool_name != expected_tool_name - { - return Err(FlowError::Internal(format!( - "nemo_guardrails returned modified tool {field} content for unexpected tool '{tool_name}'" - ))); - } - Ok(object.get(field).cloned()) -} - -fn blocking_rail_name(response: &Json) -> Option { - response - .get("guardrails") - .and_then(|guardrails| guardrails.get("log")) - .and_then(|log| log.get("activated_rails")) - .and_then(Json::as_array) - .and_then(|activated| { - activated.iter().find_map(|rail| { - let stopped = rail.get("stop").and_then(Json::as_bool) == Some(true); - let refused = - rail.get("decisions") - .and_then(Json::as_array) - .is_some_and(|decisions| { - decisions.iter().any(|decision| { - decision - .as_str() - .is_some_and(|decision| decision.starts_with("refuse ")) - }) - }); - if stopped || refused { - rail.get("name").and_then(Json::as_str).map(str::to_string) - } else { - None - } - }) - }) -} - -fn remote_mark_data( - stream: bool, - config_id: &Option, - config_ids: &[String], - status: Option, - error: Option, -) -> Json { - let mut data = Map::new(); - data.insert("stream".to_string(), Json::Bool(stream)); - if let Some(config_id) = config_id { - data.insert("config_id".to_string(), Json::String(config_id.clone())); - } - if !config_ids.is_empty() { - data.insert( - "config_ids".to_string(), - Json::Array(config_ids.iter().cloned().map(Json::String).collect()), - ); - } - if let Some(status) = status { - data.insert( - "http_status".to_string(), - Json::Number(serde_json::Number::from(status)), - ); - } - if let Some(error) = error { - data.insert("error".to_string(), Json::String(error)); - } - Json::Object(data) -} - -fn tool_remote_mark_data( - kind: RemoteCheckKind, - tool_name: &str, - config_id: &Option, - config_ids: &[String], - status: Option, - error: Option, -) -> Json { - let mut data = match remote_mark_data(false, config_id, config_ids, status, error) { - Json::Object(data) => data, - _ => unreachable!("remote_mark_data always returns an object"), - }; - data.insert( - "surface".to_string(), - Json::String(match kind { - RemoteCheckKind::Input => "tool_input".to_string(), - RemoteCheckKind::Output => "tool_output".to_string(), - }), - ); - data.insert("tool_name".to_string(), Json::String(tool_name.to_string())); - Json::Object(data) -} - -pub(super) fn register_remote_backend( - config: NeMoGuardrailsConfig, - ctx: &mut PluginRegistrationContext, -) -> PluginResult<()> { - let runtime = Arc::new(RemoteBackendRuntime::new(&config)?); - - if config.input || config.output { - let llm_execution_runtime = Arc::clone(&runtime); - let llm_execution: LlmExecutionFn = Arc::new(move |_name, request, _context, _next| { - let runtime = Arc::clone(&llm_execution_runtime); - Box::pin(async move { runtime.execute(request, false).await }) - }); - ctx.register_llm_execution_intercept("llm_remote_backend", config.priority, llm_execution)?; - - let llm_stream_runtime = Arc::clone(&runtime); - let llm_stream_execution: LlmStreamExecutionFn = - Arc::new(move |_name, request, _context, _next| { - let runtime = Arc::clone(&llm_stream_runtime); - Box::pin(async move { runtime.execute_stream(request).await }) - }); - ctx.register_llm_stream_execution_intercept( - "llm_stream_remote_backend", - config.priority, - llm_stream_execution, - )?; - } - - if config.tool_input || config.tool_output { - let tool_runtime = Arc::clone(&runtime); - let enable_tool_input = config.tool_input; - let enable_tool_output = config.tool_output; - let tool_execution: ToolExecutionFn = Arc::new(move |context, next| { - let runtime = Arc::clone(&tool_runtime); - let tool_name = context.tool_name().to_string(); - let args = context.into_args(); - Box::pin(async move { - let current_args = if enable_tool_input { - runtime.check_tool_input(&tool_name, &args).await? - } else { - args - }; - - let mut execution_result = next(current_args.clone()).await?; - if enable_tool_output { - execution_result.result = runtime - .check_tool_output(&tool_name, ¤t_args, &execution_result.result) - .await?; - } - Ok(execution_result.into()) - }) - }); - ctx.register_tool_execution_intercept( - "tool_remote_backend", - config.priority, - tool_execution, - )?; - } - - Ok(()) -} - -fn build_base_guardrails_config( - config_id: &Option, - config_ids: &[String], - request_defaults: Option<&RequestDefaultsConfig>, -) -> Map { - let mut guardrails = Map::new(); - if let Some(config_id) = config_id { - guardrails.insert("config_id".to_string(), Json::String(config_id.clone())); - } - if !config_ids.is_empty() { - guardrails.insert( - "config_ids".to_string(), - Json::Array(config_ids.iter().cloned().map(Json::String).collect()), - ); - } - if let Some(request_defaults) = request_defaults { - if let Some(context) = &request_defaults.context { - guardrails.insert("context".to_string(), context.clone()); - } - if let Some(thread_id) = &request_defaults.thread_id { - guardrails.insert("thread_id".to_string(), Json::String(thread_id.clone())); - } - if let Some(state) = &request_defaults.state { - guardrails.insert("state".to_string(), state.clone()); - } - } - guardrails -} - -fn build_llm_guardrails_config( - config_id: &Option, - config_ids: &[String], - request_defaults: Option<&RequestDefaultsConfig>, - input_enabled: bool, - output_enabled: bool, -) -> Option> { - let mut guardrails = build_base_guardrails_config(config_id, config_ids, request_defaults); - let options = build_llm_options(request_defaults, input_enabled, output_enabled); - - if !options.is_empty() { - guardrails.insert("options".to_string(), Json::Object(options)); - } - (!guardrails.is_empty()).then_some(guardrails) -} - -fn build_llm_options( - request_defaults: Option<&RequestDefaultsConfig>, - input_enabled: bool, - output_enabled: bool, -) -> Map { - let mut options = Map::new(); - if let Some(rails) = build_llm_rails_option(request_defaults, input_enabled, output_enabled) { - options.insert("rails".to_string(), Json::Object(rails)); - } - insert_llm_request_default_options(&mut options, request_defaults); - options -} - -fn build_llm_rails_option( - request_defaults: Option<&RequestDefaultsConfig>, - input_enabled: bool, - output_enabled: bool, -) -> Option> { - let mut rails = request_defaults - .and_then(|defaults| defaults.rails.as_ref()) - .map(serialize_request_rails) - .unwrap_or_default(); - - if !input_enabled { - rails.insert("input".to_string(), Json::Bool(false)); - } - if !output_enabled { - rails.insert("output".to_string(), Json::Bool(false)); - } - - (!rails.is_empty()).then_some(rails) -} - -fn insert_llm_request_default_options( - options: &mut Map, - request_defaults: Option<&RequestDefaultsConfig>, -) { - let Some(request_defaults) = request_defaults else { - return; - }; - - if let Some(llm_params) = &request_defaults.llm_params { - options.insert("llm_params".to_string(), llm_params.clone()); - } - if let Some(llm_output) = request_defaults.llm_output { - options.insert("llm_output".to_string(), Json::Bool(llm_output)); - } - if let Some(output_vars) = &request_defaults.output_vars { - options.insert("output_vars".to_string(), output_vars.clone()); - } - if let Some(log) = &request_defaults.log { - options.insert("log".to_string(), log.clone()); - } -} - -fn build_tool_check_guardrails_config( - kind: RemoteCheckKind, - config_id: &Option, - config_ids: &[String], - request_defaults: Option<&RequestDefaultsConfig>, -) -> Map { - let mut guardrails = build_base_guardrails_config(config_id, config_ids, request_defaults); - let mut options = Map::new(); - let mut rails = Map::from_iter([ - ("input".to_string(), Json::Bool(false)), - ("output".to_string(), Json::Bool(false)), - ("dialog".to_string(), Json::Bool(false)), - ("retrieval".to_string(), Json::Bool(false)), - ]); - match kind { - RemoteCheckKind::Input => { - rails.insert("tool_input".to_string(), Json::Bool(false)); - rails.insert( - "tool_output".to_string(), - configured_tool_selector(request_defaults, RemoteCheckKind::Input) - .unwrap_or(Json::Bool(true)), - ); - } - RemoteCheckKind::Output => { - rails.insert( - "tool_input".to_string(), - configured_tool_selector(request_defaults, RemoteCheckKind::Output) - .unwrap_or(Json::Bool(true)), - ); - rails.insert("tool_output".to_string(), Json::Bool(false)); - } - }; - options.insert("rails".to_string(), Json::Object(rails)); - let mut log = request_defaults - .and_then(|defaults| defaults.log.as_ref()) - .and_then(Json::as_object) - .cloned() - .unwrap_or_default(); - log.insert("activated_rails".to_string(), Json::Bool(true)); - options.insert("log".to_string(), Json::Object(log)); - guardrails.insert("options".to_string(), Json::Object(options)); - guardrails -} - -fn serialize_request_rails(rails: &RequestRailsConfig) -> Map { - serde_json::to_value(rails) - .expect("request rails config should serialize to JSON") - .as_object() - .cloned() - .expect("request rails config should serialize to a JSON object") -} - -fn configured_tool_selector( - request_defaults: Option<&RequestDefaultsConfig>, - kind: RemoteCheckKind, -) -> Option { - let rails = request_defaults.and_then(|defaults| defaults.rails.as_ref())?; - match kind { - RemoteCheckKind::Input => rails.tool_input.as_ref(), - RemoteCheckKind::Output => rails.tool_output.as_ref(), - } - .map(|selector| { - serde_json::to_value(selector).expect("tool rail selector should serialize to JSON") - }) -} - -#[cfg(test)] -#[path = "../../../tests/unit/plugins/nemo_guardrails/remote_coverage_tests.rs"] -mod coverage_tests; diff --git a/crates/core/tests/unit/plugin_tests.rs b/crates/core/tests/unit/plugin_tests.rs index 5743814918..8775316f4e 100644 --- a/crates/core/tests/unit/plugin_tests.rs +++ b/crates/core/tests/unit/plugin_tests.rs @@ -1043,6 +1043,73 @@ fn test_validate_plugin_config_honors_policy_and_duplicate_singletons() { reset_global(); } +#[test] +fn test_removed_nemo_guardrails_component_is_always_rejected() { + let _guard = lock_runtime_owner(); + reset_global(); + + for enabled in [true, false] { + for unknown_component in [ + UnsupportedBehavior::Ignore, + UnsupportedBehavior::Warn, + UnsupportedBehavior::Error, + ] { + let report = test_validate_static_plugin_config(&PluginConfig { + components: vec![PluginComponentSpec { + kind: "nemo_guardrails".into(), + enabled, + config: Default::default(), + }], + policy: ConfigPolicy { + unknown_component, + ..ConfigPolicy::default() + }, + ..PluginConfig::default() + }); + + assert_eq!(report.diagnostics.len(), 1); + let diagnostic = report + .diagnostics + .iter() + .find(|diagnostic| diagnostic.code == "plugin.removed_component") + .expect("removed Guardrails component should have a migration diagnostic"); + assert_eq!(diagnostic.level, DiagnosticLevel::Error); + assert_eq!(diagnostic.component.as_deref(), Some("nemo_guardrails")); + assert!( + diagnostic + .message + .contains("removed in NeMo Relay >=0.10.0") + ); + assert!(diagnostic.message.contains("migration-guides")); + } + } + + reset_global(); +} + +#[test] +fn test_dynamic_nemo_guardrails_component_is_not_treated_as_legacy_static_config() { + let _guard = lock_runtime_owner(); + reset_global(); + + let report = validate_plugin_config_with_static_component_count( + &PluginConfig { + components: vec![PluginComponentSpec::new("nemo_guardrails")], + ..PluginConfig::default() + }, + 0, + ); + + assert!( + report + .diagnostics + .iter() + .all(|diagnostic| diagnostic.code != "plugin.removed_component") + ); + + reset_global(); +} + #[test] fn test_validate_plugin_config_passes_top_level_policy_to_plugins() { let _guard = lock_runtime_owner(); @@ -3090,9 +3157,11 @@ fn replacing_plugin_configuration_removes_old_callbacks_and_installs_new_ones() components: vec![PluginComponentSpec::new("test.plugin")], ..PluginConfig::default() }; + let original_component_count = original.components.len(); runtime .block_on(initialize_plugins_exact_inner( original.clone(), + original_component_count, None, vec![], )) @@ -3107,6 +3176,7 @@ fn replacing_plugin_configuration_removes_old_callbacks_and_installs_new_ones() runtime .block_on(initialize_plugins_exact_inner( PluginConfig::default(), + 0, None, vec![], )) @@ -3120,7 +3190,12 @@ fn replacing_plugin_configuration_removes_old_callbacks_and_installs_new_ones() ); runtime - .block_on(initialize_plugins_exact_inner(original, None, vec![])) + .block_on(initialize_plugins_exact_inner( + original, + original_component_count, + None, + vec![], + )) .unwrap(); assert_eq!(list_runtime_registrations(None).unwrap(), before); clear_plugin_configuration_inner().result.unwrap(); @@ -3157,8 +3232,14 @@ fn failed_plugin_replacement_restores_callbacks_and_runtime_diagnostics() { original .components .push(PluginComponentSpec::new("test.plugin")); + let original_component_count = original.components.len(); let initial_report = runtime - .block_on(initialize_plugins_exact_inner(original, None, vec![])) + .block_on(initialize_plugins_exact_inner( + original, + original_component_count, + None, + vec![], + )) .unwrap(); assert!( initial_report @@ -3192,6 +3273,7 @@ fn failed_plugin_replacement_restores_callbacks_and_runtime_diagnostics() { components: vec![PluginComponentSpec::new("partial.fail.plugin")], ..PluginConfig::default() }, + 1, Some(failures.clone()), vec![], )) @@ -3247,8 +3329,14 @@ fn failed_plugin_replacement_reports_failure_to_restore_the_previous_configurati .enable_all() .build() .unwrap(); + let component_count = config.components.len(); let error = runtime - .block_on(initialize_plugins_exact_inner(config, None, vec![])) + .block_on(initialize_plugins_exact_inner( + config, + component_count, + None, + vec![], + )) .unwrap_err(); assert!(matches!(error, PluginError::RegistrationFailed(_))); assert!( @@ -3293,6 +3381,7 @@ fn replacement_teardown_failure_preserves_diagnostics_and_records_unremoved_call let error = runtime .block_on(initialize_plugins_exact_inner( PluginConfig::default(), + 0, Some(failures.clone()), vec![], )) diff --git a/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs b/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs deleted file mode 100644 index 99e578c822..0000000000 --- a/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs +++ /dev/null @@ -1,1152 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Unit tests for the built-in NeMo Guardrails plugin component contract. -#![allow(clippy::await_holding_lock)] - -use super::*; -use crate::api::runtime::NemoRelayContextState; -use std::io::{Read, Write}; -use std::net::TcpListener; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex, mpsc}; -use std::thread; -use std::time::Duration; - -use crate::api::event::Event; -use crate::api::llm::{ - LlmAttributes, LlmCallExecuteParams, LlmRequest, LlmStreamCallExecuteParams, llm_call_execute, - llm_stream_call_execute, -}; -use crate::api::runtime::global_context; -use crate::api::runtime::{ - LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, create_scope_stack, - set_thread_scope_stack, -}; -use crate::api::subscriber::{deregister_subscriber, register_subscriber}; -use crate::api::tool::{ToolCallExecuteParams, tool_call_execute}; -use crate::codec::openai_chat::{OpenAIChatCodec, OpenAIChatStreamingCodec}; -use crate::codec::streaming::StreamingCodec; -use crate::codec::traits::LlmResponseCodec; -use crate::config_editor::{EditorConfig, EditorFieldKind}; -#[cfg(feature = "schema")] -use crate::plugin::plugin_config_schema; -use crate::plugin::{ - PluginComponentSpec, PluginConfig, list_plugin_kinds, lookup_plugin, test_close_plugin_host, - test_initialize_plugin_host_exact, test_validate_static_plugin_config, -}; -use futures::StreamExt; -use serde_json::json; - -const TEST_TIMEOUT: Duration = Duration::from_secs(5); - -fn reset_runtime() { - let _ = spdlog::init_log_crate_proxy(); - log::set_max_level(log::LevelFilter::Info); - let _ = test_close_plugin_host(); - crate::shared_runtime::reset_runtime_owner_for_tests(); - let context = global_context(); - *context.write().unwrap() = NemoRelayContextState::new(); -} - -fn setup_isolated_thread() { - let stack = create_scope_stack(); - set_thread_scope_stack(stack); -} - -fn component(config: Json) -> PluginComponentSpec { - let Json::Object(config) = config else { - panic!("component config must be an object"); - }; - PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: true, - config, - } -} - -fn disabled_component(config: Json) -> PluginComponentSpec { - let Json::Object(config) = config else { - panic!("component config must be an object"); - }; - PluginComponentSpec { - kind: NEMO_GUARDRAILS_PLUGIN_KIND.to_string(), - enabled: false, - config, - } -} - -fn plugin_config(config: Json) -> PluginConfig { - PluginConfig { - version: 1, - components: vec![component(config)], - policy: Default::default(), - } -} - -fn remote_valid_config() -> Json { - json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "safety-default" - } - }) -} - -#[derive(Debug)] -struct CapturedHttpRequest { - path: String, - content_type: String, - body: Vec, -} - -fn spawn_http_responder( - listener: TcpListener, - response: Vec, - request_tx: mpsc::Sender, -) { - thread::spawn(move || { - let (mut stream, _) = listener.accept().unwrap(); - let request = read_http_request(&mut stream); - stream.write_all(&response).unwrap(); - request_tx.send(request).unwrap(); - }); -} - -fn read_http_request(stream: &mut impl Read) -> CapturedHttpRequest { - let mut bytes = Vec::new(); - let mut buf = [0_u8; 4096]; - let (header_end, content_length) = read_http_headers(stream, &mut bytes, &mut buf); - read_http_body(stream, &mut bytes, &mut buf, header_end + content_length); - - let headers_text = String::from_utf8_lossy(&bytes[..header_end]); - let request_line = headers_text.lines().next().unwrap(); - CapturedHttpRequest { - path: request_line.split_whitespace().nth(1).unwrap().to_string(), - content_type: header_value(&headers_text, "content-type") - .unwrap_or_default() - .to_string(), - body: bytes[header_end..header_end + content_length].to_vec(), - } -} - -fn read_http_headers( - stream: &mut impl Read, - bytes: &mut Vec, - buf: &mut [u8; 4096], -) -> (usize, usize) { - loop { - let read = stream.read(buf).unwrap(); - if read == 0 { - panic!("remote responder closed before receiving request"); - } - bytes.extend_from_slice(&buf[..read]); - - if let Some(header_end) = bytes.windows(4).position(|window| window == b"\r\n\r\n") { - let header_end = header_end + 4; - let headers_text = String::from_utf8_lossy(&bytes[..header_end]); - let content_length = header_value(&headers_text, "content-length") - .and_then(|value| value.parse::().ok()) - .unwrap_or(0); - return (header_end, content_length); - } - } -} - -fn read_http_body( - stream: &mut impl Read, - bytes: &mut Vec, - buf: &mut [u8; 4096], - expected_total: usize, -) { - while bytes.len() < expected_total { - let read = stream.read(buf).unwrap(); - if read == 0 { - panic!("remote responder closed before full request body"); - } - bytes.extend_from_slice(&buf[..read]); - } -} - -fn header_value<'a>(headers_text: &'a str, header_name: &str) -> Option<&'a str> { - headers_text.lines().find_map(|line| { - let (name, value) = line.split_once(':')?; - if name.eq_ignore_ascii_case(header_name) { - Some(value.trim()) - } else { - None - } - }) -} - -fn recv_captured_request(request_rx: &mpsc::Receiver) -> CapturedHttpRequest { - request_rx - .recv_timeout(TEST_TIMEOUT) - .expect("timed out waiting for captured HTTP request") -} - -fn make_chat_request(stream: bool) -> LlmRequest { - LlmRequest { - headers: serde_json::Map::new(), - content: json!({ - "model": "gpt-4o-mini", - "messages": [{"role": "user", "content": "hello"}], - "temperature": 0.2, - "stream": stream - }), - } -} - -fn capture_events(name: &str) -> Arc>> { - let events = Arc::new(Mutex::new(Vec::new())); - let sink = Arc::clone(&events); - register_subscriber( - name, - Arc::new(move |event| sink.lock().unwrap().push(event.clone())), - ) - .unwrap(); - events -} - -fn captured_events_snapshot(events: &Arc>>) -> Vec { - crate::api::subscriber::flush_subscribers().unwrap(); - events.lock().unwrap().clone() -} - -fn unused_local_endpoint() -> String { - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - drop(listener); - format!("http://{address}") -} - -#[test] -fn editor_schema_tracks_nemo_guardrails_config_types() { - let schema = NeMoGuardrailsConfig::editor_schema(); - let mode = schema.field("mode").expect("mode field"); - assert_eq!(mode.kind, EditorFieldKind::Enum); - assert_eq!(mode.enum_values, &["remote", "local"]); - - let remote = schema.field("remote").expect("remote section"); - assert_eq!(remote.kind, EditorFieldKind::Section); - assert!(remote.optional); - - let remote_schema = remote.schema().expect("remote editor schema"); - let headers = remote_schema.field("headers").expect("headers field"); - assert_eq!(headers.kind, EditorFieldKind::StringMap); - let config_ids = remote_schema.field("config_ids").expect("config_ids field"); - assert_eq!(config_ids.kind, EditorFieldKind::List); - assert_eq!( - config_ids.list_item.expect("config_ids item metadata").kind, - EditorFieldKind::String - ); - - let request_defaults = schema - .field("request_defaults") - .expect("request_defaults section"); - assert_eq!(request_defaults.kind, EditorFieldKind::Section); - assert!(request_defaults.optional); - - let request_defaults_schema = request_defaults - .schema() - .expect("request_defaults editor schema"); - let rails = request_defaults_schema.field("rails").expect("rails field"); - assert_eq!(rails.kind, EditorFieldKind::Section); - - let rails_schema = rails.schema().expect("request rails editor schema"); - let retrieval = rails_schema.field("retrieval").expect("retrieval field"); - assert_eq!(retrieval.kind, EditorFieldKind::Json); -} - -#[test] -fn default_config_and_component_conversion_cover_public_shape() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - - let defaults = NeMoGuardrailsConfig::default(); - assert_eq!(defaults.version, 1); - assert_eq!(defaults.mode, "remote"); - assert!(defaults.input); - assert!(defaults.output); - assert!(!defaults.tool_input); - assert!(!defaults.tool_output); - assert_eq!(defaults.priority, 100); - assert!(defaults.remote.is_none()); - assert!(defaults.local.is_none()); - assert!(defaults.request_defaults.is_none()); - - let remote = RemoteBackendConfig::default(); - assert_eq!(remote.timeout_millis, 3_000); - assert!(remote.headers.is_empty()); - assert!(remote.config_ids.is_empty()); - - let generic: PluginComponentSpec = ComponentSpec::new(NeMoGuardrailsConfig { - remote: Some(RemoteBackendConfig { - endpoint: Some("http://localhost:8000".into()), - config_id: Some("default".into()), - ..RemoteBackendConfig::default() - }), - ..NeMoGuardrailsConfig::default() - }) - .into(); - assert_eq!(generic.kind, NEMO_GUARDRAILS_PLUGIN_KIND); - assert!(generic.enabled); - assert_eq!(generic.config["mode"], json!("remote")); - assert_eq!(generic.config["remote"]["config_id"], json!("default")); -} - -#[cfg(feature = "schema")] -fn schema_has_property(schema: &Json, name: &str) -> bool { - schema_property(schema, name).is_some() -} - -#[cfg(feature = "schema")] -fn schema_property_has_enum(schema: &Json, name: &str, expected: &[&str]) -> bool { - schema_property(schema, name) - .and_then(|property| property.get("enum")) - .and_then(Json::as_array) - .is_some_and(|values| { - expected - .iter() - .all(|expected| values.iter().any(|value| value == *expected)) - }) -} - -#[cfg(feature = "schema")] -fn schema_property_has_default(schema: &Json, name: &str, expected: Json) -> bool { - schema_property(schema, name) - .and_then(|property| property.get("default")) - .is_some_and(|default| default == &expected) -} - -#[cfg(feature = "schema")] -fn schema_property<'a>(schema: &'a Json, name: &str) -> Option<&'a Json> { - match schema { - Json::Object(object) => { - if let Some(property) = object - .get("properties") - .and_then(Json::as_object) - .and_then(|properties| properties.get(name)) - { - return Some(property); - } - object - .values() - .find_map(|value| schema_property(value, name)) - } - Json::Array(values) => values.iter().find_map(|value| schema_property(value, name)), - _ => None, - } -} - -#[cfg(feature = "schema")] -#[test] -fn schema_contains_every_supported_nemo_guardrails_option() { - let schema = nemo_guardrails_config_schema(); - assert_eq!(schema.get("deprecated"), Some(&json!(true))); - for definition in [ - "LocalBackendConfig", - "RemoteBackendConfig", - "RequestDefaultsConfig", - "RequestRailsConfig", - "RailSelector", - ] { - assert_eq!( - schema.pointer(&format!("/definitions/{definition}/deprecated")), - Some(&json!(true)), - "schema definition `{definition}` is not marked deprecated" - ); - } - for field in [ - "version", - "mode", - "config_path", - "config_yaml", - "colang_content", - "codec", - "input", - "output", - "tool_input", - "tool_output", - "priority", - "remote", - "local", - "request_defaults", - "policy", - "endpoint", - "config_id", - "config_ids", - "headers", - "timeout_millis", - "python_module", - "python_executable", - "python_path", - "context", - "thread_id", - "state", - "rails", - "llm_params", - "llm_output", - "output_vars", - "log", - "retrieval", - "dialog", - "unknown_component", - "unknown_field", - "unsupported_value", - ] { - assert!( - schema_has_property(&schema, field), - "schema missing property `{field}`:\n{}", - serde_json::to_string_pretty(&schema).unwrap() - ); - } - assert!(schema_property_has_enum( - &schema, - "mode", - &["remote", "local"] - )); - assert!(schema_property_has_enum( - &schema, - "codec", - &[ - "openai_chat", - "openai_responses", - "anthropic_messages", - "oci_genai", - "gemini_generate_content" - ] - )); - assert!(schema_property_has_default( - &schema, - "mode", - json!("remote") - )); -} - -#[cfg(feature = "schema")] -#[test] -fn plugin_schema_contains_generic_plugin_surface() { - let schema = plugin_config_schema(); - for field in [ - "version", - "components", - "policy", - "kind", - "enabled", - "config", - ] { - assert!( - schema_has_property(&schema, field), - "plugin schema missing property `{field}`" - ); - } -} - -#[test] -fn builtin_registration_is_automatic() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - - assert!(list_plugin_kinds().contains(&NEMO_GUARDRAILS_PLUGIN_KIND.to_string())); - assert!(lookup_plugin(NEMO_GUARDRAILS_PLUGIN_KIND).is_some()); -} - -#[test] -fn configured_component_reports_deprecation_warning() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - - for component in [ - component(remote_valid_config()), - disabled_component(remote_valid_config()), - ] { - let report = test_validate_static_plugin_config(&PluginConfig { - version: 1, - components: vec![component], - policy: Default::default(), - }); - - assert!(!report.has_errors()); - let diagnostic = report - .diagnostics - .iter() - .find(|diag| diag.code == NEMO_GUARDRAILS_DEPRECATION_CODE) - .expect("configured NeMo Guardrails component should report deprecation"); - assert_eq!(diagnostic.level, DiagnosticLevel::Warning); - assert_eq!( - diagnostic.component.as_deref(), - Some(NEMO_GUARDRAILS_PLUGIN_KIND) - ); - assert!(diagnostic.message.contains(NEMO_GUARDRAILS_REMOVAL_VERSION)); - } -} - -#[test] -fn explicit_registration_helpers_are_idempotent_and_reversible() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - - assert!(register_nemo_guardrails_component().is_ok()); - assert!(register_nemo_guardrails_component().is_ok()); - assert!(deregister_nemo_guardrails_component()); - assert!(!deregister_nemo_guardrails_component()); - register_nemo_guardrails_component().unwrap(); -} - -#[test] -fn disabled_component_validates_and_initializes_without_runtime_work() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - - let config = PluginConfig { - version: 1, - components: vec![disabled_component(remote_valid_config())], - policy: Default::default(), - }; - assert!(!test_validate_static_plugin_config(&config).has_errors()); - let report = futures::executor::block_on(test_initialize_plugin_host_exact(config)).unwrap(); - assert!( - report - .diagnostics - .iter() - .any(|diagnostic| diagnostic.code == NEMO_GUARDRAILS_DEPRECATION_CODE) - ); -} - -#[test] -fn duplicate_component_is_rejected_as_singleton() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - - let config = PluginConfig { - version: 1, - components: vec![ - component(remote_valid_config()), - component(remote_valid_config()), - ], - policy: Default::default(), - }; - let report = test_validate_static_plugin_config(&config); - assert!(report.has_errors()); - assert!( - report - .diagnostics - .iter() - .any(|diag| diag.code == "plugin.duplicate_component") - ); -} - -#[test] -fn invalid_shapes_and_values_are_reported() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - assert_invalid_shape_and_mode(); - assert_invalid_local_config(); - assert_invalid_remote_identity_and_codec(); - assert_remote_tool_surface_validation(); - assert_empty_and_mixed_config_values(); - assert_request_defaults_validation(); -} - -fn assert_invalid_shape_and_mode() { - let invalid_shape = test_validate_static_plugin_config(&plugin_config(json!({ - "version": "one", - }))); - assert!(invalid_shape.has_errors()); - assert!( - invalid_shape - .diagnostics - .iter() - .any(|diag| diag.code == "nemo_guardrails.invalid_plugin_config") - ); - - let unsupported_version_and_mode = test_validate_static_plugin_config(&plugin_config(json!({ - "version": 2, - "mode": "hybrid", - "codec": "openai_chat", - "remote": {"endpoint": "http://localhost:8000", "config_id": "default"} - }))); - assert!(unsupported_version_and_mode.has_errors()); - assert!( - unsupported_version_and_mode - .diagnostics - .iter() - .any( - |diag| diag.code == "nemo_guardrails.unsupported_config_version" - && diag.field.as_deref() == Some("version") - ) - ); - assert!( - unsupported_version_and_mode - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("mode") - && diag.message.contains("mode must be 'remote' or 'local'")) - ); -} - -fn assert_invalid_local_config() { - let local_missing_source = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "local", - "codec": "openai_chat", - }))); - assert!(local_missing_source.has_errors()); - assert!(local_missing_source.diagnostics.iter().any(|diag| { - diag.message - .contains("exactly one of config_path or config_yaml is required in local mode") - })); - - let local_bad_colang = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "local", - "config_path": "./rails", - "colang_content": "define flow x", - "codec": "openai_chat", - }))); - assert!(local_bad_colang.has_errors()); - assert!( - local_bad_colang - .diagnostics - .iter() - .any(|diag| diag.message.contains("colang_content can only be used")) - ); - - let local_rejects_remote_section = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "local", - "config_yaml": "rails:\n input: []\n", - "codec": "openai_chat", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - } - }))); - assert!(local_rejects_remote_section.has_errors()); - assert!( - local_rejects_remote_section - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("remote") - && diag.message.contains("cannot be used when mode is 'local'")) - ); -} - -fn assert_invalid_remote_identity_and_codec() { - let remote_missing_identity = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": {"endpoint": "http://localhost:8000"}, - }))); - assert!(remote_missing_identity.has_errors()); - assert!(remote_missing_identity.diagnostics.iter().any(|diag| { - diag.message - .contains("remote mode requires remote.config_id or remote.config_ids") - })); - - let remote_conflicting_ids = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "one", - "config_ids": ["two"] - }, - }))); - assert!(remote_conflicting_ids.has_errors()); - assert!(remote_conflicting_ids.diagnostics.iter().any(|diag| { - diag.message - .contains("remote.config_id and remote.config_ids cannot be used together") - })); - - let missing_codec = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - } - }))); - assert!(missing_codec.has_errors()); - assert!( - missing_codec - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("codec")) - ); - - let bad_codec = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "openai_agents", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - } - }))); - assert!(bad_codec.has_errors()); - assert!(bad_codec.diagnostics.iter().any(|diag| { - diag.message.contains("codec must be one of:") - && [ - "openai_chat", - "openai_responses", - "anthropic_messages", - "oci_genai", - "gemini_generate_content", - ] - .iter() - .all(|name| diag.message.contains(name)) - })); - - let unsupported_remote_codec = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "openai_responses", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - } - }))); - assert!(unsupported_remote_codec.has_errors()); - assert!(unsupported_remote_codec.diagnostics.iter().any(|diag| { - diag.message - .contains("remote mode currently supports only codec = 'openai_chat'") - })); - - let unsupported_remote_anthropic_codec = - test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "anthropic_messages", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - } - }))); - assert!(unsupported_remote_anthropic_codec.has_errors()); - assert!( - unsupported_remote_anthropic_codec - .diagnostics - .iter() - .any(|diag| { - diag.message - .contains("remote mode currently supports only codec = 'openai_chat'") - }) - ); - - let unsupported_remote_oci_codec = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "oci_genai", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - } - }))); - assert!(unsupported_remote_oci_codec.has_errors()); - assert!(unsupported_remote_oci_codec.diagnostics.iter().any(|diag| { - diag.message - .contains("remote mode currently supports only codec = 'openai_chat'") - })); - - let unsupported_remote_gemini_codec = - test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "gemini_generate_content", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - } - }))); - assert!(unsupported_remote_gemini_codec.has_errors()); - assert!( - unsupported_remote_gemini_codec - .diagnostics - .iter() - .any(|diag| { - diag.message - .contains("remote mode currently supports only codec = 'openai_chat'") - }) - ); -} - -fn assert_remote_tool_surface_validation() { - let unsupported_remote_tool_input = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "tool_input": true, - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - } - }))); - assert!(unsupported_remote_tool_input.has_errors()); - assert!( - unsupported_remote_tool_input - .diagnostics - .iter() - .any(|diag| { - diag.field.as_deref() == Some("tool_input") - && diag - .message - .contains("does not currently support managed tool_input") - }) - ); - - let supported_remote_tool_output = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "tool_output": true, - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - } - }))); - assert!(!supported_remote_tool_output.has_errors()); -} - -fn assert_empty_and_mixed_config_values() { - let remote_empty_fields = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": "", - "config_id": "", - "config_ids": ["default", ""] - } - }))); - assert!(remote_empty_fields.has_errors()); - assert!( - remote_empty_fields - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("remote.endpoint")) - ); - assert!( - remote_empty_fields - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("remote.config_id")) - ); - assert!( - remote_empty_fields - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("remote.config_ids[1]")) - ); - - let remote_local_mix = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "config_path": "./rails", - "codec": "openai_chat", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - }, - "local": {"python_module": "nemoguardrails"} - }))); - assert!(remote_local_mix.has_errors()); - assert!( - remote_local_mix - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("local")) - ); - assert!(remote_local_mix.diagnostics.iter().any(|diag| { - diag.message - .contains("remote mode uses remote config identity") - })); - - let no_surfaces = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "local", - "config_path": "./rails", - "input": false, - "output": false, - "tool_input": false, - "tool_output": false - }))); - assert!(no_surfaces.has_errors()); - assert!( - no_surfaces - .diagnostics - .iter() - .any(|diag| diag.message.contains("at least one Guardrails surface")) - ); - - let local_empty_fields = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "local", - "config_path": "", - "config_yaml": "", - "colang_content": "", - "codec": "openai_chat", - "local": {"python_module": "", "python_executable": "", "python_path": ""} - }))); - assert!(local_empty_fields.has_errors()); - assert!( - local_empty_fields - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("config_path")) - ); - assert!( - local_empty_fields - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("config_yaml")) - ); - assert!( - local_empty_fields - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("colang_content")) - ); - assert!( - local_empty_fields - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("local.python_module")) - ); - assert!( - local_empty_fields - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("local.python_executable")) - ); - assert!( - local_empty_fields - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("local.python_path")) - ); -} - -fn assert_request_defaults_validation() { - let local_request_defaults = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "local", - "codec": "openai_chat", - "config_path": "./rails", - "request_defaults": { - "context": {"tenant": "demo"} - } - }))); - assert!(local_request_defaults.has_errors()); - assert!(local_request_defaults.diagnostics.iter().any(|diag| { - diag.field.as_deref() == Some("request_defaults") - && diag - .message - .contains("local mode does not currently support request_defaults") - })); - - let invalid_request_defaults = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - }, - "request_defaults": { - "context": true, - "thread_id": " ", - "state": {"foo": "bar"}, - "llm_params": [], - "log": "verbose", - "output_vars": ["answer", "", 7], - "rails": { - "retrieval": [""] - } - } - }))); - assert!(invalid_request_defaults.has_errors()); - assert!( - invalid_request_defaults - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("request_defaults.context")) - ); - assert!( - invalid_request_defaults - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("request_defaults.thread_id")) - ); - assert!(invalid_request_defaults.diagnostics.iter().any(|diag| { - diag.message - .contains("request_defaults.thread_id must not be empty") - })); - assert!( - invalid_request_defaults - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("request_defaults.state")) - ); - assert!(invalid_request_defaults.diagnostics.iter().any(|diag| { - diag.message - .contains("request_defaults.state must be empty or contain only 'events' or 'state'") - })); - assert!( - invalid_request_defaults - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("request_defaults.llm_params")) - ); - assert!( - invalid_request_defaults - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("request_defaults.log")) - ); - assert!( - invalid_request_defaults - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("request_defaults.output_vars[1]")) - ); - assert!( - invalid_request_defaults - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("request_defaults.output_vars[2]")) - ); - assert!( - invalid_request_defaults - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("request_defaults.rails.retrieval[0]")) - ); - - let invalid_request_output_vars_shape = - test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - }, - "request_defaults": { - "thread_id": "short", - "output_vars": 7 - } - }))); - assert!(invalid_request_output_vars_shape.has_errors()); - assert!( - invalid_request_output_vars_shape - .diagnostics - .iter() - .any( - |diag| diag.field.as_deref() == Some("request_defaults.thread_id") - && diag - .message - .contains("request_defaults.thread_id must be at least 16 characters long") - ) - ); - assert!( - invalid_request_output_vars_shape - .diagnostics - .iter() - .any(|diag| diag.field.as_deref() == Some("request_defaults.output_vars")) - ); - - let valid_bool_output_vars = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - }, - "request_defaults": { - "output_vars": true - } - }))); - assert!(!valid_bool_output_vars.has_errors()); -} - -#[test] -fn unknown_fields_follow_policy() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - - let warn_report = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": {"endpoint": "http://localhost:8000", "config_id": "default"}, - "bogus": true - }))); - assert!( - warn_report - .diagnostics - .iter() - .any(|diag| diag.code == "nemo_guardrails.unknown_field") - ); - - let nested_warn_report = test_validate_static_plugin_config(&plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": {"endpoint": "http://localhost:8000", "config_id": "default"}, - "request_defaults": { - "rails": { - "bogus": true - } - } - }))); - assert!( - nested_warn_report - .diagnostics - .iter() - .any(|diag| diag.component.as_deref() == Some("request_defaults.rails")) - ); - - let ignored = test_validate_static_plugin_config(&plugin_config(json!({ - "policy": {"unknown_field": "ignore", "unsupported_value": "ignore"}, - "mode": "remote", - "codec": "openai_chat", - "remote": {"endpoint": "http://localhost:8000", "config_id": "default"}, - "bogus": true - }))); - assert!(!ignored.has_errors()); - assert_eq!(ignored.diagnostics.len(), 1); - assert_eq!( - ignored.diagnostics[0].code, - NEMO_GUARDRAILS_DEPRECATION_CODE - ); -} - -#[test] -fn enabled_unknown_mode_initialization_fails_fast_when_policy_ignores_validation() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - - let error = - futures::executor::block_on(test_initialize_plugin_host_exact(plugin_config(json!({ - "policy": {"unsupported_value": "ignore"}, - "mode": "hybrid", - "codec": "openai_chat", - "remote": { - "endpoint": "http://localhost:8000", - "config_id": "default" - } - })))) - .unwrap_err(); - - match error { - crate::plugin::PluginError::InvalidConfig(message) => { - assert!(message.contains("unsupported NeMo Guardrails mode 'hybrid'")); - } - other => panic!("unexpected error: {other}"), - } -} - -#[path = "remote_tests.rs"] -mod remote_tests; diff --git a/crates/core/tests/unit/plugins/nemo_guardrails/local_python_tests.rs b/crates/core/tests/unit/plugins/nemo_guardrails/local_python_tests.rs deleted file mode 100644 index 20a0ed6bb8..0000000000 --- a/crates/core/tests/unit/plugins/nemo_guardrails/local_python_tests.rs +++ /dev/null @@ -1,1476 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -#![allow(clippy::await_holding_lock)] // Runtime isolation requires serial async plugin tests. - -#[cfg(unix)] -use std::fs; -#[cfg(unix)] -use std::os::unix::fs::PermissionsExt; -#[cfg(unix)] -use std::path::{Path, PathBuf}; -#[cfg(unix)] -use std::process::Command; -#[cfg(unix)] -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; - -use serde_json::json; - -use super::*; -#[cfg(unix)] -use crate::api::llm::{LlmAttributes, LlmCallExecuteParams, LlmRequest, llm_call_execute}; -#[cfg(unix)] -use crate::api::runtime::{ - LlmExecutionNextFn, NemoRelayContextState, ThreadScopeStackBinding, capture_thread_scope_stack, - create_scope_stack, global_context, restore_thread_scope_stack, set_thread_scope_stack, -}; -#[cfg(unix)] -use crate::api::tool::{ToolCallExecuteParams, tool_call_execute}; -#[cfg(unix)] -use crate::codec::openai_chat::OpenAIChatCodec; -#[cfg(unix)] -use crate::codec::traits::LlmResponseCodec; -#[cfg(unix)] -use crate::plugin::{ - PluginComponentSpec, PluginConfig, test_close_plugin_host, test_initialize_plugin_host_exact, -}; -use crate::plugins::nemo_guardrails::component::LocalBackendConfig; - -#[cfg(unix)] -static NEXT_FIXTURE_ID: AtomicUsize = AtomicUsize::new(1); -static PYTHON_EXECUTABLE_ENV_MUTEX: Mutex<()> = Mutex::new(()); - -struct EnvVarGuard { - name: &'static str, - value: Option, -} - -impl EnvVarGuard { - fn set(name: &'static str, value: &str) -> Self { - let old_value = std::env::var_os(name); - unsafe { - std::env::set_var(name, value); - } - Self { - name, - value: old_value, - } - } - - fn remove(name: &'static str) -> Self { - let old_value = std::env::var_os(name); - unsafe { - std::env::remove_var(name); - } - Self { - name, - value: old_value, - } - } -} - -impl Drop for EnvVarGuard { - fn drop(&mut self) { - unsafe { - match &self.value { - Some(value) => std::env::set_var(self.name, value), - None => std::env::remove_var(self.name), - } - } - } -} - -#[test] -fn python_executable_prefers_config_over_environment() { - let _env_guard = PYTHON_EXECUTABLE_ENV_MUTEX.lock().unwrap(); - let _nemo_python = EnvVarGuard::set(PYTHON_EXECUTABLE_ENV, "env-python"); - let _pyo3_python = EnvVarGuard::set(PYO3_PYTHON_ENV, "pyo3-python"); - let _uv_python = EnvVarGuard::set(UV_PYTHON_ENV, "uv-python"); - - let config = NeMoGuardrailsConfig { - local: Some(LocalBackendConfig { - python_executable: Some("configured-python".to_string()), - ..LocalBackendConfig::default() - }), - ..NeMoGuardrailsConfig::default() - }; - - assert_eq!(python_executable(&config), "configured-python"); -} - -#[test] -fn python_executable_uses_python_environment_before_default() { - let _env_guard = PYTHON_EXECUTABLE_ENV_MUTEX.lock().unwrap(); - let _nemo_python = EnvVarGuard::remove(PYTHON_EXECUTABLE_ENV); - let _pyo3_python = EnvVarGuard::set(PYO3_PYTHON_ENV, "pyo3-python"); - let _uv_python = EnvVarGuard::set(UV_PYTHON_ENV, "uv-python"); - - assert_eq!( - python_executable(&NeMoGuardrailsConfig::default()), - "pyo3-python" - ); -} - -#[test] -fn local_worker_start_reports_an_unavailable_python_executable() { - let config = NeMoGuardrailsConfig { - local: Some(LocalBackendConfig { - python_executable: Some("nemo-relay-python-that-does-not-exist".to_string()), - ..LocalBackendConfig::default() - }), - ..NeMoGuardrailsConfig::default() - }; - - let error = LocalGuardrailsWorker::start(&config) - .err() - .expect("an unavailable executable should fail worker startup"); - assert!( - error - .to_string() - .contains("failed to start NeMo Guardrails local Python worker") - ); -} - -#[test] -fn worker_python_path_prepends_configured_path_to_inherited_pythonpath() { - let configured = std::path::PathBuf::from("fake-guardrails"); - let stdlib = std::path::PathBuf::from("stdlib"); - let platstdlib = std::path::PathBuf::from("platstdlib"); - let configured_path = std::env::join_paths([configured.clone()]).unwrap(); - let inherited_path = std::env::join_paths([stdlib.clone(), platstdlib.clone()]).unwrap(); - - let merged = merge_python_path(&configured_path, Some(&inherited_path)).unwrap(); - - assert_eq!( - std::env::split_paths(&merged).collect::>(), - vec![configured, stdlib, platstdlib] - ); -} - -#[cfg(unix)] -struct FakeGuardrails { - root: PathBuf, - module_name: String, - python: PathBuf, -} - -#[cfg(unix)] -impl FakeGuardrails { - fn new(version: &str) -> Self { - let _ = spdlog::init_log_crate_proxy(); - log::set_max_level(log::LevelFilter::Info); - let id = NEXT_FIXTURE_ID.fetch_add(1, Ordering::Relaxed); - let module_name = format!("fake_guardrails_{id}"); - let root = std::env::temp_dir().join(format!( - "nemo_relay_fake_guardrails_{}_{}", - std::process::id(), - id - )); - let package = root.join(&module_name); - fs::create_dir_all(package.join("rails/llm")).unwrap(); - fs::write(package.join("rails/__init__.py"), "").unwrap(); - fs::write(package.join("rails/llm/__init__.py"), "").unwrap(); - fs::write(package.join("rails/llm/options.py"), fake_options_module()).unwrap(); - fs::write(package.join("__init__.py"), fake_root_module(version)).unwrap(); - - let python = root.join("python-wrapper"); - fs::write( - &python, - format!( - "#!/bin/sh\nPYTHONPATH='{}' exec python3 \"$@\"\n", - shell_single_quote(&root) - ), - ) - .unwrap(); - let mut permissions = fs::metadata(&python).unwrap().permissions(); - permissions.set_mode(0o755); - fs::set_permissions(&python, permissions).unwrap(); - - Self { - root, - module_name, - python, - } - } - - fn config(&self) -> NeMoGuardrailsConfig { - NeMoGuardrailsConfig { - mode: "local".to_string(), - codec: Some("openai_chat".to_string()), - config_yaml: Some("models: []".to_string()), - colang_content: Some("define flow noop\n pass".to_string()), - local: Some(LocalBackendConfig { - python_module: Some(self.module_name.clone()), - python_executable: Some(self.python.to_string_lossy().into_owned()), - python_path: None, - }), - ..NeMoGuardrailsConfig::default() - } - } -} - -#[cfg(unix)] -impl Drop for FakeGuardrails { - fn drop(&mut self) { - let _ = fs::remove_dir_all(&self.root); - } -} - -#[cfg(unix)] -fn python3_available() -> bool { - Command::new("python3") - .arg("--version") - .output() - .map(|output| output.status.success()) - .unwrap_or(false) -} - -#[test] -#[cfg(unix)] -fn shutdown_handles_an_already_reaped_worker_process() { - let _ = spdlog::init_log_crate_proxy(); - log::set_max_level(log::LevelFilter::Info); - let mut child = Command::new("true").spawn().unwrap(); - child.wait().unwrap(); - let worker = LocalGuardrailsWorker { - writer: Mutex::new(None), - child: Mutex::new(child), - waiters: Arc::new(Mutex::new(HashMap::new())), - stream_events: Arc::new(Mutex::new(HashMap::new())), - next_id: AtomicU64::new(0), - shutdown_started: AtomicBool::new(false), - }; - - worker.shutdown(); -} - -#[tokio::test] -#[cfg(unix)] -async fn request_timeout_stops_an_unresponsive_worker() { - let _ = spdlog::init_log_crate_proxy(); - log::set_max_level(log::LevelFilter::Info); - let child = Command::new("sleep").arg("60").spawn().unwrap(); - let (sender, receiver) = std_mpsc::channel::(); - let waiters = Arc::new(Mutex::new(HashMap::new())); - let stream_events = Arc::new(Mutex::new(HashMap::new())); - let closed_waiters = Arc::clone(&waiters); - let closed_stream_events = Arc::clone(&stream_events); - let handle = thread::spawn(move || { - while receiver.recv().is_ok() { - thread::sleep(Duration::from_millis(50)); - notify_worker_closed( - &closed_waiters, - &closed_stream_events, - "test worker closed".into(), - ); - } - }); - let worker = LocalGuardrailsWorker { - writer: Mutex::new(Some(WorkerCommandWriter { - sender, - error: Arc::new(Mutex::new(None)), - handle: Some(handle), - })), - child: Mutex::new(child), - waiters, - stream_events, - next_id: AtomicU64::new(0), - shutdown_started: AtomicBool::new(false), - }; - - let error = worker - .request_with_timeout(json!({"command": "never-reply"}), Duration::from_millis(10)) - .await; - assert!(error.unwrap_err().to_string().contains("timed out")); -} - -#[cfg(unix)] -fn shell_single_quote(path: &Path) -> String { - path.to_string_lossy().replace('\'', "'\\''") -} - -#[cfg(unix)] -fn fake_options_module() -> &'static str { - r#" -class RailType: - INPUT = "input" - OUTPUT = "output" - -class RailStatus: - BLOCKED = "blocked" - MODIFIED = "modified" - PASSED = "passed" -"# -} - -#[cfg(unix)] -fn fake_root_module(version: &str) -> String { - format!( - r#" -import json -import types -from .rails.llm.options import RailStatus - -__version__ = {version:?} - -class Result: - def __init__(self, status, content=None, rail=None): - self.status = status - self.content = content - self.rail = rail - -class RailsConfig: - @staticmethod - def from_content(*, colang_content=None, yaml_content=None): - stream_first = "stream_first_false" not in (yaml_content or "") - flows = [] if "no_stream" in (yaml_content or "") else ["self check output"] - return types.SimpleNamespace( - yaml=yaml_content, - colang=colang_content, - rails=types.SimpleNamespace( - output=types.SimpleNamespace( - flows=flows, - streaming=types.SimpleNamespace(enabled=True, stream_first=stream_first), - ) - ) - ) - - @staticmethod - def from_path(path): - return types.SimpleNamespace( - path=path, - rails=types.SimpleNamespace( - output=types.SimpleNamespace( - flows=["self check output"], - streaming=types.SimpleNamespace(enabled=True, stream_first=True), - ) - ) - ) - -class LLMRails: - def __init__(self, config): - self.config = config - - async def check_async(self, messages, rail_types=None): - content = " ".join(str(message.get("content", "")) for message in messages) - if "block" in content: - return Result(RailStatus.BLOCKED, "", "policy") - if "modify-tool" in content: - return Result(RailStatus.MODIFIED, '{{"arguments":{{"safe":true}},"result":{{"ok":true}}}}') - if "modify" in content: - return Result(RailStatus.MODIFIED, "rewritten") - return Result(RailStatus.PASSED, "") - - async def stream_async(self, *, messages=None, generator=None, include_metadata=False): - async for text in generator: - if "stream-block" in text: - yield json.dumps({{"error": {{"type": "guardrails_violation", "message": "blocked stream"}}}}) - return - yield json.dumps({{"ok": True}}) -"# - ) -} - -#[cfg(unix)] -#[tokio::test(flavor = "current_thread")] -async fn bridge_checks_pass_block_and_modify_outcomes() { - if !python3_available() { - return; - } - - let fixture = FakeGuardrails::new("0.22.0"); - let bridge = LocalGuardrailsBridge::new(&fixture.config()).unwrap(); - - assert!(matches!( - bridge - .check( - vec![json!({"role": "user", "content": "hello"})], - LocalRailKind::Input, - ) - .await - .unwrap(), - LocalCheckOutcome::Passed - )); - - match bridge - .check( - vec![json!({"role": "user", "content": "block this"})], - LocalRailKind::Input, - ) - .await - .unwrap() - { - LocalCheckOutcome::Blocked { rail } => assert_eq!(rail.as_deref(), Some("policy")), - _ => panic!("expected blocked outcome"), - } - - match bridge - .check( - vec![json!({"role": "user", "content": "modify this"})], - LocalRailKind::Input, - ) - .await - .unwrap() - { - LocalCheckOutcome::Modified { content } => assert_eq!(content, "rewritten"), - _ => panic!("expected modified outcome"), - } -} - -#[cfg(unix)] -#[test] -fn bridge_rejects_unsupported_guardrails_version() { - if !python3_available() { - return; - } - - let fixture = FakeGuardrails::new("0.21.0"); - let error = match LocalGuardrailsBridge::new(&fixture.config()) { - Ok(_) => panic!("expected unsupported version error"), - Err(error) => error, - }; - assert!(error.to_string().contains("nemoguardrails==0.22.0")); -} - -#[cfg(unix)] -#[tokio::test(flavor = "current_thread")] -async fn streaming_support_rejects_stream_first_false() { - if !python3_available() { - return; - } - - let fixture = FakeGuardrails::new("0.22.0"); - let mut config = fixture.config(); - config.config_yaml = Some("stream_first_false".to_string()); - let bridge = LocalGuardrailsBridge::new(&config).unwrap(); - - assert!(bridge.has_streaming_output_rails().await.unwrap()); - let error = bridge - .ensure_streaming_output_supported() - .await - .unwrap_err(); - assert!(error.to_string().contains("stream_first = true")); -} - -#[cfg(unix)] -#[tokio::test(flavor = "current_thread")] -async fn stream_monitor_records_blocked_message() { - if !python3_available() { - return; - } - - let fixture = FakeGuardrails::new("0.22.0"); - let bridge = LocalGuardrailsBridge::new(&fixture.config()).unwrap(); - let (text_tx, text_rx) = mpsc::channel(8); - let blocked = Arc::new(Mutex::new(None)); - let monitor = bridge - .spawn_stream_monitor( - vec![json!({"role": "user", "content": "hello"})], - text_rx, - Arc::clone(&blocked), - ) - .unwrap(); - - text_tx - .send(Some("stream-block".to_string())) - .await - .unwrap(); - text_tx.send(None).await.unwrap(); - monitor.await.unwrap().unwrap(); - - assert_eq!(blocked.lock().unwrap().as_deref(), Some("blocked stream")); -} - -#[tokio::test(flavor = "current_thread")] -async fn guarded_provider_stream_reports_block_after_forwarded_chunks() { - let provider_stream = LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ - "choices": [{"delta": {"content": "blocked"}}], - }))])); - let (text_tx, mut text_rx) = mpsc::channel::>(8); - let (chunk_tx, mut chunk_rx) = mpsc::channel(8); - let blocked = Arc::new(Mutex::new(None)); - let monitor_blocked = Arc::clone(&blocked); - let monitor = tokio::spawn(async move { - while let Some(item) = text_rx.recv().await { - match item { - Some(text) if text.contains("blocked") => { - *monitor_blocked.lock().unwrap() = Some("blocked stream".to_string()); - } - Some(_) => {} - None => break, - } - } - Ok(()) - }); - let (_cancel, cancel_rx) = tokio::sync::watch::channel(false); - let (closed, _closed_rx) = tokio::sync::watch::channel(None); - - forward_guarded_provider_stream( - provider_stream, - LocalGuardrailsCodec::OpenAIChat, - text_tx, - chunk_tx, - monitor, - blocked, - cancel_rx, - closed, - ) - .await; - - let chunk = chunk_rx.recv().await.unwrap().unwrap(); - assert_eq!( - chunk, - json!({ - "choices": [{"delta": {"content": "blocked"}}], - }) - ); - - let error = chunk_rx.recv().await.unwrap().unwrap_err(); - assert!( - error.to_string().contains("blocked stream"), - "unexpected error: {error}" - ); - assert!(chunk_rx.recv().await.is_none()); -} - -#[tokio::test] -async fn stream_monitor_errors_are_forwarded_to_the_provider_stream() { - async fn panicking_monitor() -> FlowResult<()> { - panic!("monitor panicked"); - } - - let blocked = Arc::new(Mutex::new(None)); - - for monitor in [ - tokio::spawn(async { Err(FlowError::Internal("monitor failed".into())) }), - tokio::spawn(panicking_monitor()), - ] { - let (chunk_tx, mut chunk_rx) = mpsc::channel(1); - assert!(send_stream_monitor_error(monitor, &chunk_tx, &blocked).await); - assert!(chunk_rx.recv().await.unwrap().is_err()); - } - - let (chunk_tx, mut chunk_rx) = mpsc::channel(1); - *blocked.lock().unwrap() = Some("blocked output".into()); - assert!(send_stream_monitor_error(tokio::spawn(async { Ok(()) }), &chunk_tx, &blocked).await); - assert!( - chunk_rx - .recv() - .await - .unwrap() - .unwrap_err() - .to_string() - .contains("blocked output") - ); - - *blocked.lock().unwrap() = None; - assert!(!send_stream_monitor_error(tokio::spawn(async { Ok(()) }), &chunk_tx, &blocked).await); -} - -#[tokio::test] -async fn guarded_provider_stream_forwards_provider_and_channel_failures() { - let provider_error = FlowError::Internal("provider failed".into()); - let provider_stream = LlmJsonStream::new(tokio_stream::iter(vec![Err(provider_error)])); - let (text_tx, mut text_rx) = mpsc::channel(2); - let (chunk_tx, mut chunk_rx) = mpsc::channel(1); - let (_cancel_tx, cancel_rx) = watch::channel(false); - let (closed_tx, closed_rx) = watch::channel(None); - forward_guarded_provider_stream( - provider_stream, - LocalGuardrailsCodec::OpenAIChat, - text_tx, - chunk_tx, - tokio::spawn(async { Ok(()) }), - Arc::new(Mutex::new(None)), - cancel_rx, - closed_tx, - ) - .await; - assert!(chunk_rx.recv().await.unwrap().is_err()); - assert_eq!(text_rx.recv().await, Some(None)); - assert!(closed_rx.borrow().as_ref().unwrap().is_ok()); - - let provider_stream = LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ - "choices": [{"delta": {"content": "hello"}}] - }))])); - let (text_tx, text_rx) = mpsc::channel(1); - drop(text_rx); - let (chunk_tx, mut chunk_rx) = mpsc::channel(1); - let (_cancel_tx, cancel_rx) = watch::channel(false); - let (closed_tx, _closed_rx) = watch::channel(None); - forward_guarded_provider_stream( - provider_stream, - LocalGuardrailsCodec::OpenAIChat, - text_tx, - chunk_tx, - tokio::spawn(async { Err(FlowError::Internal("monitor closed".into())) }), - Arc::new(Mutex::new(None)), - cancel_rx, - closed_tx, - ) - .await; - assert!(chunk_rx.recv().await.unwrap().is_err()); -} - -#[tokio::test] -async fn guarded_provider_stream_handles_preblocked_dropped_and_cancelled_consumers() { - let stream_chunk = || { - LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({ - "choices": [{"delta": {"content": "hello"}}] - }))])) - }; - - let (text_tx, _text_rx) = mpsc::channel(3); - let (chunk_tx, mut chunk_rx) = mpsc::channel(2); - let (_cancel_tx, cancel_rx) = watch::channel(false); - let (closed_tx, _closed_rx) = watch::channel(None); - forward_guarded_provider_stream( - stream_chunk(), - LocalGuardrailsCodec::OpenAIChat, - text_tx, - chunk_tx, - tokio::spawn(async { Ok(()) }), - Arc::new(Mutex::new(Some("already blocked".into()))), - cancel_rx, - closed_tx, - ) - .await; - assert!(chunk_rx.recv().await.unwrap().is_err()); - - let (text_tx, _text_rx) = mpsc::channel(3); - let (chunk_tx, chunk_rx) = mpsc::channel(1); - drop(chunk_rx); - let (_cancel_tx, cancel_rx) = watch::channel(false); - let (closed_tx, _closed_rx) = watch::channel(None); - forward_guarded_provider_stream( - stream_chunk(), - LocalGuardrailsCodec::OpenAIChat, - text_tx, - chunk_tx, - tokio::spawn(async { Ok(()) }), - Arc::new(Mutex::new(None)), - cancel_rx, - closed_tx, - ) - .await; - - let (text_tx, mut text_rx) = mpsc::channel(1); - let (chunk_tx, _chunk_rx) = mpsc::channel(1); - let (cancel_tx, cancel_rx) = watch::channel(false); - cancel_tx.send_replace(true); - let (closed_tx, _closed_rx) = watch::channel(None); - forward_guarded_provider_stream( - stream_chunk(), - LocalGuardrailsCodec::OpenAIChat, - text_tx, - chunk_tx, - tokio::spawn(async { std::future::pending::>().await }), - Arc::new(Mutex::new(None)), - cancel_rx, - closed_tx, - ) - .await; - assert_eq!(text_rx.recv().await, Some(None)); -} - -#[test] -fn local_codec_and_rewrite_helpers_cover_all_provider_surfaces() { - for (codec, surface) in [ - ( - LocalGuardrailsCodec::OpenAIChat, - ProviderSurface::OpenAIChat, - ), - ( - LocalGuardrailsCodec::OpenAIResponses, - ProviderSurface::OpenAIResponses, - ), - ( - LocalGuardrailsCodec::AnthropicMessages, - ProviderSurface::AnthropicMessages, - ), - ( - LocalGuardrailsCodec::GeminiGenerateContent, - ProviderSurface::GeminiGenerateContent, - ), - ] { - assert_eq!(codec.provider_surface(), surface); - assert_eq!( - LocalGuardrailsCodec::from_provider_surface(surface).provider_surface(), - surface - ); - } - - let mut config = NeMoGuardrailsConfig { - input: false, - output: false, - ..Default::default() - }; - assert!(resolve_codec(&config).unwrap().is_none()); - config.input = true; - assert!(resolve_codec(&config).is_err()); - config.codec = Some("unsupported".into()); - assert!(resolve_codec(&config).is_err()); - - let mut annotated = AnnotatedLlmRequest { - messages: vec![Message::Assistant { - content: None, - tool_calls: None, - name: None, - }], - ..Default::default() - }; - replace_last_role_content(&mut annotated, "assistant", "rewritten".into()).unwrap(); - assert!(matches!( - &annotated.messages[0], - Message::Assistant { - content: Some(MessageContent::Text(content)), - .. - } if content == "rewritten" - )); - assert!(replace_last_role_content(&mut annotated, "user", "missing".into()).is_err()); - assert!(modified_tool_payload("[]", "arguments").is_err()); -} - -#[test] -fn parse_check_result_rejects_unknown_status() { - assert!(matches!( - parse_check_result(json!({"status": "passed"})).unwrap(), - LocalCheckOutcome::Passed - )); - - let error = match parse_check_result(json!({"status": "surprising"})) { - Ok(_) => panic!("expected unknown status to fail"), - Err(error) => error, - }; - assert!( - error - .to_string() - .contains("unexpected worker check status: surprising"), - "unexpected error: {error}" - ); - assert!(parse_check_result(json!({"status": 7})).is_err()); -} - -#[test] -fn worker_envelope_helpers_cover_delivery_shutdown_and_default_results() { - assert!(set_request_id(&mut Json::Null, "1").is_err()); - let mut payload = json!({"command": "check"}); - set_request_id(&mut payload, "request-1").unwrap(); - assert_eq!(payload["id"], json!("request-1")); - - let waiters = Arc::new(Mutex::new(HashMap::new())); - let stream_events = Arc::new(Mutex::new(HashMap::new())); - let (waiter_tx, waiter_rx) = std_mpsc::channel(); - waiters.lock().unwrap().insert("unary".into(), waiter_tx); - dispatch_worker_envelope( - &waiters, - &stream_events, - WorkerEnvelope { - id: "unary".into(), - ok: true, - result: Some(json!({"ok": true})), - error: None, - event: None, - message: None, - }, - ); - assert!(waiter_rx.recv().unwrap().ok); - - let (stream_tx, mut stream_rx) = mpsc::unbounded_channel(); - stream_events - .lock() - .unwrap() - .insert("stream".into(), stream_tx); - dispatch_worker_envelope( - &waiters, - &stream_events, - WorkerEnvelope { - id: "stream".into(), - ok: true, - result: None, - error: None, - event: Some("done".into()), - message: None, - }, - ); - assert_eq!(stream_rx.try_recv().unwrap().event.as_deref(), Some("done")); - - let (waiter_tx, waiter_rx) = std_mpsc::channel(); - waiters.lock().unwrap().insert("closed".into(), waiter_tx); - let (stream_tx, mut stream_rx) = mpsc::unbounded_channel(); - stream_events - .lock() - .unwrap() - .insert("closed-stream".into(), stream_tx); - notify_worker_closed(&waiters, &stream_events, "worker gone".into()); - assert_eq!( - waiter_rx.recv().unwrap().error.as_deref(), - Some("worker gone") - ); - assert_eq!( - stream_rx.try_recv().unwrap().error.as_deref(), - Some("worker gone") - ); - - assert_eq!( - worker_result(WorkerEnvelope { - id: "ok".into(), - ok: true, - result: None, - error: None, - event: None, - message: None, - }) - .unwrap(), - Json::Null - ); - assert!( - worker_result(WorkerEnvelope { - id: "error".into(), - ok: false, - result: None, - error: None, - event: None, - message: None, - }) - .unwrap_err() - .to_string() - .contains("worker failed") - ); -} - -#[test] -fn worker_command_writer_reports_stored_and_closed_channel_errors() { - let (sender, receiver) = std_mpsc::channel(); - let writer = WorkerCommandWriter { - sender, - error: Arc::new(Mutex::new(Some("broken pipe".into()))), - handle: None, - }; - assert!( - writer - .send("ignored".into()) - .unwrap_err() - .to_string() - .contains("broken pipe") - ); - drop(receiver); - - let (sender, receiver) = std_mpsc::channel(); - drop(receiver); - let writer = WorkerCommandWriter { - sender, - error: Arc::new(Mutex::new(None)), - handle: None, - }; - assert!( - writer - .send("ignored".into()) - .unwrap_err() - .to_string() - .contains("channel closed") - ); -} - -#[cfg(unix)] -#[test] -fn worker_reader_handles_blank_valid_invalid_and_eof_lines() { - { - let worker = monitor_test_worker(); - let (sender, receiver) = std_mpsc::channel(); - worker.waiters.lock().unwrap().insert("ok".into(), sender); - let mut valid_source = Command::new("sh") - .arg("-c") - .arg("printf '\n{\"id\":\"ok\",\"ok\":true}\n'") - .stdout(Stdio::piped()) - .spawn() - .unwrap(); - worker.spawn_reader(valid_source.stdout.take().unwrap()); - assert!(receiver.recv_timeout(Duration::from_secs(1)).unwrap().ok); - valid_source.wait().unwrap(); - } - - let worker = monitor_test_worker(); - let (sender, receiver) = std_mpsc::channel(); - worker - .waiters - .lock() - .unwrap() - .insert("invalid".into(), sender); - let mut invalid_source = Command::new("sh") - .arg("-c") - .arg("printf 'not-json\n'") - .stdout(Stdio::piped()) - .spawn() - .unwrap(); - worker.spawn_reader(invalid_source.stdout.take().unwrap()); - assert!( - receiver - .recv_timeout(Duration::from_secs(1)) - .unwrap() - .error - .unwrap() - .contains("invalid worker response") - ); - invalid_source.wait().unwrap(); -} - -#[cfg(unix)] -#[test] -fn closed_worker_writer_cleans_up_unary_and_stream_registrations() { - let worker = monitor_test_worker(); - let mut request = json!({"command": "check"}); - assert!(worker.send_request(&mut request).is_err()); - assert!(worker.waiters.lock().unwrap().is_empty()); - - assert!(worker.start_stream(vec![json!({"role": "user"})]).is_err()); - assert!(worker.stream_events.lock().unwrap().is_empty()); -} - -#[tokio::test] -async fn guarded_stream_close_reports_an_early_cleanup_exit() { - let (_chunk_tx, chunk_rx) = mpsc::channel(1); - let (cancel, _cancel_rx) = watch::channel(false); - let (closed_tx, closed) = watch::channel(None); - drop(closed_tx); - let mut stream = GuardedProviderStream { - receiver: ReceiverStream::new(chunk_rx), - cancel, - closed, - }; - assert!( - Pin::new(&mut stream) - .close() - .await - .unwrap_err() - .to_string() - .contains("cleanup task ended early") - ); -} - -#[cfg(unix)] -fn monitor_test_worker() -> Arc { - Arc::new(LocalGuardrailsWorker { - writer: Mutex::new(None), - child: Mutex::new(Command::new("sleep").arg("60").spawn().unwrap()), - waiters: Arc::new(Mutex::new(HashMap::new())), - stream_events: Arc::new(Mutex::new(HashMap::new())), - next_id: AtomicU64::new(0), - shutdown_started: AtomicBool::new(false), - }) -} - -#[cfg(unix)] -async fn run_monitor_event(event: Option) -> (FlowResult<()>, Option) { - let worker = monitor_test_worker(); - let (text_tx, text_rx) = mpsc::channel(1); - let (event_tx, event_rx) = mpsc::unbounded_channel(); - let blocked = Arc::new(Mutex::new(None)); - if let Some(event) = event { - event_tx.send(event).unwrap(); - } - drop(event_tx); - let result = monitor_guardrails_stream( - worker, - "stream-id".into(), - text_rx, - event_rx, - Arc::clone(&blocked), - ) - .await; - drop(text_tx); - let message = blocked.lock().unwrap().clone(); - (result, message) -} - -#[cfg(unix)] -#[tokio::test] -async fn stream_monitor_handles_terminal_worker_event_variants() { - let envelope = |ok, event: &str, error: Option<&str>, message: Option<&str>| WorkerEnvelope { - id: "stream-id".into(), - ok, - result: None, - error: error.map(str::to_string), - event: Some(event.into()), - message: message.map(str::to_string), - }; - - let (result, blocked) = run_monitor_event(Some(envelope( - true, - "blocked", - None, - Some("policy blocked output"), - ))) - .await; - assert!(result.is_ok()); - assert_eq!(blocked.as_deref(), Some("policy blocked output")); - - assert!( - run_monitor_event(Some(envelope(true, "done", None, None))) - .await - .0 - .is_ok() - ); - assert!( - run_monitor_event(Some(envelope(false, "error", None, None))) - .await - .0 - .unwrap_err() - .to_string() - .contains("worker stream failed") - ); - assert!( - run_monitor_event(Some(envelope(true, "unexpected", None, None))) - .await - .0 - .unwrap_err() - .to_string() - .contains("unknown stream event") - ); - assert!( - run_monitor_event(None) - .await - .0 - .unwrap_err() - .to_string() - .contains("closed unexpectedly") - ); -} - -#[test] -fn modified_tool_payload_rejects_malformed_content() { - let error = modified_tool_payload("not-json", "arguments").unwrap_err(); - assert!( - error - .to_string() - .contains("modified tool arguments content that is not valid JSON") - ); - - let error = modified_tool_payload(r#"{"tool_name":"demo"}"#, "result").unwrap_err(); - assert!( - error - .to_string() - .contains("modified tool result content without a 'result' field") - ); -} - -#[test] -fn stream_text_extraction_handles_supported_codecs() { - assert_eq!( - extract_stream_text( - LocalGuardrailsCodec::OpenAIChat, - &json!({"choices": [{"delta": {"content": "hel"}}, {"delta": {"content": "lo"}}]}) - ), - Some("hello".to_string()) - ); - assert_eq!( - extract_stream_text( - LocalGuardrailsCodec::OpenAIResponses, - &json!({"type": "response.output_text.delta", "delta": "hello"}) - ), - Some("hello".to_string()) - ); - assert_eq!( - extract_stream_text( - LocalGuardrailsCodec::AnthropicMessages, - &json!({"type": "content_block_delta", "delta": {"type": "text_delta", "text": "hello"}}) - ), - Some("hello".to_string()) - ); - // OCI GENERIC: bare choice delta with a top-level `message`. - assert_eq!( - extract_stream_text( - LocalGuardrailsCodec::OCIGenAI, - &json!({"index": 0, "message": {"content": [{"type": "TEXT", "text": "hello"}]}}) - ), - Some("hello".to_string()) - ); - // OCI GENERIC: `choices`-wrapped deltas, optionally inside `chatResponse`. - assert_eq!( - extract_stream_text( - LocalGuardrailsCodec::OCIGenAI, - &json!({"chatResponse": {"choices": [ - {"index": 0, "message": {"content": [{"type": "TEXT", "text": "hel"}]}}, - {"index": 1, "message": {"content": [{"type": "TEXT", "text": "lo"}]}} - ]}}) - ), - Some("hello".to_string()) - ); - // OCI COHERE: bare text fragment. - assert_eq!( - extract_stream_text(LocalGuardrailsCodec::OCIGenAI, &json!({"text": "hello"})), - Some("hello".to_string()) - ); - // Gemini: visible text parts reach the guardrail worker. - assert_eq!( - extract_stream_text( - LocalGuardrailsCodec::GeminiGenerateContent, - &json!({"candidates": [{"content": {"parts": [{"text": "visible"}]}, "index": 0}]}) - ), - Some("visible".to_string()) - ); - for (codec, chunk) in [ - (LocalGuardrailsCodec::OpenAIChat, Json::Null), - ( - LocalGuardrailsCodec::OpenAIResponses, - json!({"type": "response.completed", "delta": "ignored"}), - ), - ( - LocalGuardrailsCodec::AnthropicMessages, - json!({"type": "message_delta", "delta": {"type": "text_delta", "text": "ignored"}}), - ), - ( - LocalGuardrailsCodec::AnthropicMessages, - json!({"type": "content_block_delta", "delta": {"type": "input_json_delta"}}), - ), - (LocalGuardrailsCodec::GeminiGenerateContent, Json::Null), - // A tool-call-only OCI delta carries no user-visible text and must not - // reach the guardrail worker. - ( - LocalGuardrailsCodec::OCIGenAI, - json!({"index": 0, "message": {"content": [], "toolCalls": [{"arguments": "{"}]}}), - ), - // The terminal COHERE event repeats the full response text alongside - // finishReason; forwarding it would double the rail input. - ( - LocalGuardrailsCodec::OCIGenAI, - json!({"apiFormat": "COHERE", "text": "hello!", "finishReason": "COMPLETE"}), - ), - ( - LocalGuardrailsCodec::OCIGenAI, - json!({"chatResponse": {"apiFormat": "COHERE", "text": "hello!", "finishReason": "COMPLETE"}}), - ), - ] { - assert_eq!(extract_stream_text(codec, &chunk), None); - } -} - -#[test] -fn stream_text_extraction_oci_cohere_stream_is_not_doubled() { - // Live-shaped COHERE stream: incremental deltas, then a terminal event - // repeating the complete text. The rails must see the text exactly once. - let stream = [ - json!({"apiFormat": "COHERE", "text": "hello"}), - json!({"apiFormat": "COHERE", "text": "!"}), - json!({"apiFormat": "COHERE", "text": "hello!", "finishReason": "COMPLETE"}), - ]; - let forwarded: String = stream - .iter() - .filter_map(|chunk| extract_stream_text(LocalGuardrailsCodec::OCIGenAI, chunk)) - .collect(); - assert_eq!(forwarded, "hello!"); -} - -#[test] -fn stream_text_extraction_gemini_skips_thought_parts() { - // A thought chunk (thought: true) must NOT reach the guardrail worker. - assert_eq!( - extract_stream_text( - LocalGuardrailsCodec::GeminiGenerateContent, - &json!({"candidates": [{"content": {"parts": [{"thought": true, "text": "internal reasoning"}]}, "index": 0}]}) - ), - None, - "thought parts must not be forwarded to the guardrail worker" - ); - // A chunk with both a thought part and a visible part: only the visible text is forwarded. - assert_eq!( - extract_stream_text( - LocalGuardrailsCodec::GeminiGenerateContent, - &json!({"candidates": [{"content": {"parts": [ - {"thought": true, "text": "reasoning"}, - {"text": "answer"} - ]}, "index": 0}]}) - ), - Some("answer".to_string()), - "only non-thought text must reach the guardrail worker" - ); -} - -#[cfg(unix)] -async fn install_local_plugin(config: &NeMoGuardrailsConfig) { - let component_config = serde_json::to_value(config) - .unwrap() - .as_object() - .unwrap() - .clone(); - test_initialize_plugin_host_exact(PluginConfig { - version: 1, - components: vec![PluginComponentSpec { - kind: crate::plugins::nemo_guardrails::component::NEMO_GUARDRAILS_PLUGIN_KIND - .to_string(), - enabled: true, - config: component_config, - }], - policy: Default::default(), - }) - .await - .unwrap(); -} - -#[cfg(unix)] -struct PluginRuntimeResetGuard { - previous_scope_stack: ThreadScopeStackBinding, -} - -#[cfg(unix)] -impl Drop for PluginRuntimeResetGuard { - fn drop(&mut self) { - let _ = test_close_plugin_host(); - crate::shared_runtime::reset_runtime_owner_for_tests(); - *global_context().write().unwrap() = NemoRelayContextState::new(); - restore_thread_scope_stack(self.previous_scope_stack.clone()); - } -} - -#[cfg(unix)] -fn reset_plugin_runtime() -> PluginRuntimeResetGuard { - let previous_scope_stack = capture_thread_scope_stack(); - let _ = test_close_plugin_host(); - crate::shared_runtime::reset_runtime_owner_for_tests(); - *global_context().write().unwrap() = NemoRelayContextState::new(); - set_thread_scope_stack(create_scope_stack()); - PluginRuntimeResetGuard { - previous_scope_stack, - } -} - -#[cfg(unix)] -#[tokio::test(flavor = "current_thread")] -async fn registered_local_backend_rewrites_llm_requests_and_tool_payloads() { - if !python3_available() { - return; - } - - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - let _runtime_guard = reset_plugin_runtime(); - - let fixture = FakeGuardrails::new("0.22.0"); - let mut config = fixture.config(); - config.input = true; - config.output = true; - config.tool_input = true; - config.tool_output = true; - install_local_plugin(&config).await; - - let observed_request = Arc::new(Mutex::new(None)); - let observed_callback_request = Arc::clone(&observed_request); - let callback: LlmExecutionNextFn = Arc::new(move |request| { - *observed_callback_request.lock().unwrap() = Some(request); - Box::pin(async move { - Ok(json!({ - "choices": [{"message": {"role": "assistant", "content": "provider answer"}}] - })) - }) - }); - let request = LlmRequest { - headers: Default::default(), - content: json!({ - "model": "gpt-4o-mini", - "messages": [{"role": "user", "content": "modify this request"}] - }), - }; - let response = llm_call_execute( - LlmCallExecuteParams::builder() - .name("openai") - .request(request) - .func(callback) - .attributes(LlmAttributes::empty()) - .response_codec(Arc::new(OpenAIChatCodec) as Arc) - .build(), - ) - .await - .unwrap(); - assert_eq!( - observed_request.lock().unwrap().as_ref().unwrap().content["messages"][0]["content"], - json!("rewritten") - ); - assert_eq!( - response["choices"][0]["message"]["content"], - json!("provider answer") - ); - - let tool_result = tool_call_execute( - ToolCallExecuteParams::builder() - .name("lookup") - .args(json!({"modify-tool": true})) - .func(Arc::new(|args| { - Box::pin(async move { - assert_eq!(args, json!({"safe": true})); - Ok(crate::api::tool::ToolExecutionResult::annotated( - json!({"original": true}), - json!({"source": "provider"}), - )) - }) - })) - .build(), - ) - .await - .unwrap(); - assert_eq!(tool_result.result, json!({"original": true})); - assert_eq!(tool_result.annotation, Some(json!({"source": "provider"}))); -} - -#[cfg(unix)] -#[tokio::test(flavor = "current_thread")] -async fn registered_local_backend_rejects_blocked_llm_and_tool_inputs() { - if !python3_available() { - return; - } - - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - let _runtime_guard = reset_plugin_runtime(); - - let fixture = FakeGuardrails::new("0.22.0"); - let mut config = fixture.config(); - config.input = true; - config.tool_input = true; - install_local_plugin(&config).await; - - let llm_callback_called = Arc::new(AtomicBool::new(false)); - let llm_callback_marker = Arc::clone(&llm_callback_called); - let llm_error = llm_call_execute( - LlmCallExecuteParams::builder() - .name("openai") - .request(LlmRequest { - headers: Default::default(), - content: json!({ - "model": "gpt-4o-mini", - "messages": [{"role": "user", "content": "block this request"}] - }), - }) - .func(Arc::new(move |_| { - llm_callback_marker.store(true, Ordering::SeqCst); - Box::pin(async { Ok(json!({})) }) - })) - .attributes(LlmAttributes::empty()) - .response_codec(Arc::new(OpenAIChatCodec) as Arc) - .build(), - ) - .await - .unwrap_err(); - assert!( - llm_error - .to_string() - .contains("input rail blocked the LLM call") - ); - assert!(!llm_callback_called.load(Ordering::SeqCst)); - - let tool_callback_called = Arc::new(AtomicBool::new(false)); - let tool_callback_marker = Arc::clone(&tool_callback_called); - let tool_error = tool_call_execute( - ToolCallExecuteParams::builder() - .name("lookup") - .args(json!({"block": true})) - .func(Arc::new(move |_| { - tool_callback_marker.store(true, Ordering::SeqCst); - Box::pin(async { Ok(json!({}).into()) }) - })) - .build(), - ) - .await - .unwrap_err(); - assert!( - tool_error - .to_string() - .contains("tool_input rail blocked the tool call") - ); - assert!(!tool_callback_called.load(Ordering::SeqCst)); -} - -#[cfg(unix)] -#[tokio::test(flavor = "current_thread")] -async fn registered_local_backend_guards_streams_and_rewrites_provider_requests() { - use crate::api::llm::{LlmStreamCallExecuteParams, llm_stream_call_execute}; - if !python3_available() { - return; - } - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - let _runtime_guard = reset_plugin_runtime(); - let fixture = FakeGuardrails::new("0.22.0"); - let mut config = fixture.config(); - config.input = true; - config.output = true; - install_local_plugin(&config).await; - - for (text, blocked) in [("safe answer", false), ("stream-block", true)] { - let chunk = json!({"choices": [{"delta": {"content": text}}]}); - let provider_chunk = chunk.clone(); - let mut stream = llm_stream_call_execute( - LlmStreamCallExecuteParams::builder() - .name("openai") - .request(LlmRequest { - headers: Default::default(), - content: json!({"messages": [{"role": "user", "content": "modify this request"}]}), - }) - .func(Arc::new(move |request| { - assert_eq!(request.content["messages"][0]["content"], "rewritten"); - let chunk = provider_chunk.clone(); - Box::pin(async move { Ok(LlmJsonStream::new(tokio_stream::iter([Ok(chunk)]))) }) - })) - .collector(Box::new(|_| Ok(()))) - .finalizer(Box::new(|| Json::Null)) - .build(), - ) - .await - .unwrap(); - assert_eq!(stream.next().await.unwrap().unwrap(), chunk); - if blocked { - let error = stream.next().await.unwrap().unwrap_err(); - assert!(error.to_string().contains("blocked stream"), "{error}"); - } - assert!(stream.next().await.is_none()); - } -} - -#[cfg(unix)] -#[tokio::test(flavor = "current_thread")] -async fn local_stream_runtime_respects_disabled_rails_and_rejects_unsupported_streaming() { - if !python3_available() { - return; - } - let fixture = FakeGuardrails::new("0.22.0"); - for (yaml, enable_output, rejected) in [ - ("models: []", false, false), - ("no_stream", true, false), - ("stream_first_false", true, true), - ] { - let mut config = fixture.config(); - config.config_yaml = Some(yaml.into()); - let runtime = LocalGuardrailsRuntime::new(&config).unwrap(); - let request = LlmRequest { - headers: Default::default(), - content: json!({"messages": [{"role": "user", "content": "hello"}]}), - }; - let next = Arc::new(|_: LlmRequest| { - Box::pin(async { - Ok(LlmJsonStream::new(tokio_stream::iter([Ok(json!({ - "choices": [{"delta": {"content": "stream-block"}}] - }))]))) - }) as Pin> + Send>> - }); - let result = runtime - .execute_llm_stream(request, next, false, enable_output) - .await; - if rejected { - let error = match result { - Ok(_) => panic!("unsupported streaming was accepted"), - Err(error) => error, - }; - assert!(error.to_string().contains("stream_first = true"), "{error}"); - } else { - let mut stream = result.unwrap(); - assert_eq!( - stream.next().await.unwrap().unwrap()["choices"][0]["delta"]["content"], - "stream-block" - ); - assert!(stream.next().await.is_none()); - } - } -} diff --git a/crates/core/tests/unit/plugins/nemo_guardrails/remote_coverage_tests.rs b/crates/core/tests/unit/plugins/nemo_guardrails/remote_coverage_tests.rs deleted file mode 100644 index c53caf7cea..0000000000 --- a/crates/core/tests/unit/plugins/nemo_guardrails/remote_coverage_tests.rs +++ /dev/null @@ -1,1225 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Focused remote runtime coverage tests for the NeMo Guardrails plugin component. - -use std::collections::HashMap; -use std::io::{Read, Write}; -use std::net::TcpListener; -use std::sync::Arc; -use std::time::Duration; - -use super::*; -use crate::plugins::nemo_guardrails::component::{RailSelector, RemoteBackendConfig}; -use tokio_stream::StreamExt; - -fn runtime_config(remote: RemoteBackendConfig) -> NeMoGuardrailsConfig { - let _ = spdlog::init_log_crate_proxy(); - log::set_max_level(log::LevelFilter::Info); - NeMoGuardrailsConfig { - remote: Some(remote), - ..NeMoGuardrailsConfig::default() - } -} - -fn valid_remote() -> RemoteBackendConfig { - RemoteBackendConfig { - endpoint: Some("http://127.0.0.1:1/base/".to_string()), - config_id: Some("default".to_string()), - ..RemoteBackendConfig::default() - } -} - -fn valid_runtime() -> RemoteBackendRuntime { - RemoteBackendRuntime::new(&runtime_config(valid_remote())).unwrap() -} - -fn runtime_with_endpoint(endpoint: String) -> RemoteBackendRuntime { - RemoteBackendRuntime::new(&runtime_config(RemoteBackendConfig { - endpoint: Some(endpoint), - timeout_millis: 5_000, - ..valid_remote() - })) - .unwrap() -} - -fn simple_chat_request() -> LlmRequest { - LlmRequest { - headers: Map::new(), - content: json!({ - "model": "gpt-4o-mini", - "messages": [{"role": "user", "content": "hello"}], - }), - } -} - -fn spawn_disconnecting_endpoint() -> String { - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - std::thread::spawn(move || { - let _ = listener.accept(); - }); - format!("http://{address}") -} - -fn spawn_json_response(response: Json) -> String { - spawn_http_response("200 OK", "application/json", response.to_string()) -} - -fn spawn_http_response( - status: &'static str, - content_type: &'static str, - body: impl Into>, -) -> String { - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let body = body.into(); - std::thread::spawn(move || { - let (mut stream, _) = listener.accept().unwrap(); - stream - .set_read_timeout(Some(Duration::from_secs(2))) - .unwrap(); - read_http_request(&mut stream); - write!( - stream, - "HTTP/1.1 {status}\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", - body.len() - ) - .unwrap(); - stream.write_all(&body).unwrap(); - }); - format!("http://{address}") -} - -fn spawn_truncated_http_response(status: &'static str, content_type: &'static str) -> String { - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - std::thread::spawn(move || { - let (mut stream, _) = listener.accept().unwrap(); - stream - .set_read_timeout(Some(Duration::from_secs(2))) - .unwrap(); - read_http_request(&mut stream); - write!( - stream, - "HTTP/1.1 {status}\r\nContent-Type: {content_type}\r\nContent-Length: 64\r\nConnection: close\r\n\r\npartial" - ) - .unwrap(); - }); - format!("http://{address}") -} - -fn read_http_request(stream: &mut std::net::TcpStream) { - let mut request = Vec::new(); - let mut buffer = [0; 1024]; - loop { - match stream.read(&mut buffer) { - Ok(0) => break, - Ok(n) => { - request.extend_from_slice(&buffer[..n]); - if http_request_body_complete(&request) { - break; - } - } - Err(error) - if matches!( - error.kind(), - std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut - ) => - { - break; - } - Err(error) => panic!("failed to read local HTTP request: {error}"), - } - } -} - -fn http_request_body_complete(request: &[u8]) -> bool { - let Some(header_end) = request.windows(4).position(|window| window == b"\r\n\r\n") else { - return false; - }; - let header_end = header_end + 4; - let headers = String::from_utf8_lossy(&request[..header_end]); - let content_length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().ok()) - .flatten() - }) - .unwrap_or(0); - request.len() >= header_end + content_length -} - -fn assert_flow_error_contains(result: crate::error::Result, expected: &str) { - let error = match result { - Ok(_) => panic!("expected FlowError"), - Err(error) => error, - }; - assert!( - error.to_string().contains(expected), - "expected '{error}' to contain '{expected}'" - ); -} - -fn expect_plugin_error_contains(result: PluginResult, expected: &str) { - let error = match result { - Ok(_) => panic!("expected PluginError"), - Err(error) => error, - }; - assert!( - error.to_string().contains(expected), - "expected '{error}' to contain '{expected}'" - ); -} - -#[test] -fn remote_runtime_new_reports_missing_and_invalid_config() { - expect_plugin_error_contains( - RemoteBackendRuntime::new(&NeMoGuardrailsConfig::default()), - "remote config is required", - ); - - expect_plugin_error_contains( - RemoteBackendRuntime::new(&runtime_config(RemoteBackendConfig::default())), - "remote.endpoint is required", - ); - - let mut headers = HashMap::new(); - headers.insert("bad header".to_string(), "value".to_string()); - expect_plugin_error_contains( - RemoteBackendRuntime::new(&runtime_config(RemoteBackendConfig { - headers, - ..valid_remote() - })), - "remote.headers contains invalid header name", - ); - - let mut headers = HashMap::new(); - headers.insert("x-valid".to_string(), "bad\r\nvalue".to_string()); - expect_plugin_error_contains( - RemoteBackendRuntime::new(&runtime_config(RemoteBackendConfig { - headers, - ..valid_remote() - })), - "remote.headers[x-valid] has an invalid value", - ); -} - -#[test] -fn request_body_and_guardrails_config_helpers_cover_defaults() { - let runtime = valid_runtime(); - assert_eq!( - runtime.chat_completions_url(), - "http://127.0.0.1:1/base/v1/chat/completions" - ); - - let invalid_request = LlmRequest { - headers: Map::new(), - content: Json::Null, - }; - assert_flow_error_contains( - runtime.build_request_body(&invalid_request, false), - "request content is not an object", - ); - assert_flow_error_contains( - runtime.build_request_body( - &LlmRequest { - headers: Map::new(), - content: json!({ - "model": "gpt-4o-mini", - "messages": [], - "tools": [{"type": "function", "function": {"name": "lookup"}}] - }), - }, - false, - ), - "does not support OpenAI tool definitions", - ); - runtime.record_access_status(reqwest::StatusCode::OK); - assert_eq!(runtime.access_state.load(Ordering::Acquire), 2); - runtime.record_access_status(reqwest::StatusCode::OK); - assert_eq!(runtime.access_state.load(Ordering::Acquire), 2); - - let defaults = RequestDefaultsConfig { - context: Some(json!({"tenant": "test"})), - thread_id: Some("thread-1234567890".to_string()), - state: Some(json!({"events": []})), - rails: Some(RequestRailsConfig { - input: Some(RailSelector::Enabled(true)), - output: Some(RailSelector::Enabled(true)), - retrieval: Some(RailSelector::Named(vec!["kb".to_string()])), - dialog: Some(true), - tool_input: Some(RailSelector::Named(vec!["tool-in".to_string()])), - tool_output: Some(RailSelector::Named(vec!["tool-out".to_string()])), - }), - llm_params: Some(json!({"temperature": 0.1})), - llm_output: Some(true), - output_vars: Some(json!(["answer"])), - log: Some(json!({"activated_rails": false, "details": true})), - }; - - let llm_guardrails = build_llm_guardrails_config( - &Some("primary".to_string()), - &["fallback".to_string()], - Some(&defaults), - false, - true, - ) - .expect("guardrails config"); - assert_eq!(llm_guardrails["config_id"], json!("primary")); - assert_eq!(llm_guardrails["config_ids"], json!(["fallback"])); - assert_eq!(llm_guardrails["context"], json!({"tenant": "test"})); - assert_eq!(llm_guardrails["thread_id"], json!("thread-1234567890")); - assert_eq!( - llm_guardrails["options"]["rails"]["input"], - Json::Bool(false) - ); - assert_eq!( - llm_guardrails["options"]["rails"]["retrieval"], - json!(["kb"]) - ); - assert_eq!( - llm_guardrails["options"]["llm_params"], - json!({"temperature": 0.1}) - ); - assert_eq!(llm_guardrails["options"]["output_vars"], json!(["answer"])); - assert_eq!( - build_llm_guardrails_config(&None, &[], None, true, true), - None - ); - - let tool_input = - build_tool_check_guardrails_config(RemoteCheckKind::Input, &None, &[], Some(&defaults)); - assert_eq!( - tool_input["options"]["rails"]["tool_output"], - json!(["tool-in"]) - ); - assert_eq!( - tool_input["options"]["log"]["activated_rails"], - Json::Bool(true) - ); - - let tool_output = - build_tool_check_guardrails_config(RemoteCheckKind::Output, &None, &[], Some(&defaults)); - assert_eq!( - tool_output["options"]["rails"]["tool_input"], - json!(["tool-out"]) - ); -} - -#[test] -fn named_rail_selector_combinations_are_preserved_for_llm_and_tool_checks() { - let defaults = RequestDefaultsConfig { - rails: Some(RequestRailsConfig { - input: Some(RailSelector::Named(vec![ - "input-a".to_string(), - "input-b".to_string(), - ])), - output: Some(RailSelector::Named(vec!["output-a".to_string()])), - retrieval: Some(RailSelector::Enabled(false)), - dialog: Some(false), - tool_input: Some(RailSelector::Enabled(false)), - tool_output: Some(RailSelector::Named(vec![ - "tool-output-a".to_string(), - "tool-output-b".to_string(), - ])), - }), - ..RequestDefaultsConfig::default() - }; - - let llm_guardrails = build_llm_guardrails_config( - &None, - &["named-a".to_string(), "named-b".to_string()], - Some(&defaults), - true, - true, - ) - .expect("guardrails config"); - assert_eq!(llm_guardrails["config_ids"], json!(["named-a", "named-b"])); - assert_eq!( - llm_guardrails["options"]["rails"]["input"], - json!(["input-a", "input-b"]) - ); - assert_eq!( - llm_guardrails["options"]["rails"]["output"], - json!(["output-a"]) - ); - assert_eq!( - llm_guardrails["options"]["rails"]["retrieval"], - Json::Bool(false) - ); - assert_eq!( - llm_guardrails["options"]["rails"]["dialog"], - Json::Bool(false) - ); - - let tool_input = - build_tool_check_guardrails_config(RemoteCheckKind::Input, &None, &[], Some(&defaults)); - assert_eq!( - tool_input["options"]["rails"]["tool_input"], - Json::Bool(false) - ); - assert_eq!( - tool_input["options"]["rails"]["tool_output"], - Json::Bool(false) - ); - - let tool_output = - build_tool_check_guardrails_config(RemoteCheckKind::Output, &None, &[], Some(&defaults)); - assert_eq!( - tool_output["options"]["rails"]["tool_input"], - json!(["tool-output-a", "tool-output-b"]) - ); - assert_eq!( - tool_output["options"]["rails"]["tool_output"], - Json::Bool(false) - ); -} - -#[test] -fn request_body_rejects_tool_definitions_and_sets_stream_flag() { - let runtime = valid_runtime(); - let with_tools = LlmRequest { - headers: Map::new(), - content: json!({ - "model": "gpt-4o-mini", - "messages": [{"role": "user", "content": "hello"}], - "tools": [{ - "type": "function", - "function": {"name": "search", "parameters": {"type": "object"}} - }], - }), - }; - assert_flow_error_contains( - runtime.build_request_body(&with_tools, false), - "does not support OpenAI tool definitions", - ); - - let with_tool_choice = LlmRequest { - headers: Map::new(), - content: json!({ - "model": "gpt-4o-mini", - "messages": [{"role": "user", "content": "hello"}], - "tool_choice": "auto", - }), - }; - assert_flow_error_contains( - runtime.build_request_body(&with_tool_choice, false), - "does not support OpenAI tool definitions", - ); - - let body = runtime - .build_request_body(&simple_chat_request(), true) - .expect("valid request body"); - assert_eq!(body["stream"], json!(true)); - assert_eq!(body["guardrails"]["config_id"], json!("default")); -} - -#[test] -fn tool_message_helpers_build_guardrails_compatible_chat_payloads() { - let args = json!({"city": "Phoenix"}); - let result = json!({"forecast": "sunny"}); - - let input_messages = tool_input_messages("weather_lookup", &args); - assert_eq!( - input_messages[0]["content"], - json!("Run the tool 'weather_lookup' and validate the result.") - ); - assert_eq!( - input_messages[1]["tool_calls"][0]["id"], - json!("nemo_guardrails_weather_lookup_call") - ); - assert_eq!( - input_messages[1]["tool_calls"][0]["function"]["arguments"], - json!("{\"city\":\"Phoenix\"}") - ); - - let output_messages = tool_output_messages("weather_lookup", &args, &result); - assert_eq!(output_messages[2]["role"], json!("tool")); - assert_eq!( - output_messages[2]["content"], - json!("{\"forecast\":\"sunny\"}") - ); -} - -#[test] -fn modified_tool_argument_parsing_covers_success_and_error_shapes() { - let response = json!({ - "choices": [{ - "message": { - "tool_calls": [{ - "function": { - "name": "weather_lookup", - "arguments": "{\"city\":\"Paris\"}" - } - }] - } - }] - }); - assert_eq!( - modified_tool_arguments(&response, "weather_lookup").unwrap(), - Some(json!({"city": "Paris"})) - ); - - assert_flow_error_contains( - modified_tool_arguments(&json!({"choices": []}), "weather_lookup"), - "did not contain choices[0].message", - ); - assert_flow_error_contains( - modified_tool_arguments( - &json!({"choices": [{"message": {"tool_calls": [{}]}}]}), - "weather_lookup", - ), - "without a function payload", - ); - assert_flow_error_contains( - modified_tool_arguments( - &json!({"choices": [{"message": {"tool_calls": [{"function": {}}]}}]}), - "weather_lookup", - ), - "without a function name", - ); - assert_flow_error_contains( - modified_tool_arguments( - &json!({"choices": [{"message": {"tool_calls": [{"function": {"name": "other"}}]}}]}), - "weather_lookup", - ), - "unexpected tool 'other'", - ); - assert_flow_error_contains( - modified_tool_arguments( - &json!({"choices": [{"message": {"tool_calls": [{"function": {"name": "weather_lookup"}}]}}]}), - "weather_lookup", - ), - "without function.arguments", - ); - assert_flow_error_contains( - modified_tool_arguments( - &json!({"choices": [{"message": {"tool_calls": [{"function": {"name": "weather_lookup", "arguments": "not json"}}]}}]}), - "weather_lookup", - ), - "not valid JSON", - ); - - let legacy = json!({ - "choices": [{ - "message": { - "content": "{\"tool_name\":\"weather_lookup\",\"arguments\":{\"city\":\"Berlin\"}}" - } - }] - }); - assert_eq!( - modified_tool_arguments(&legacy, "weather_lookup").unwrap(), - Some(json!({"city": "Berlin"})) - ); - assert_eq!( - modified_tool_arguments( - &json!({"choices": [{"message": {"content": "not json"}}]}), - "weather_lookup", - ) - .unwrap(), - None - ); - assert_flow_error_contains( - modified_tool_arguments( - &json!({"choices": [{"message": {"content": "{\"tool_name\":\"other\",\"arguments\":{}}"}}]}), - "weather_lookup", - ), - "unexpected tool 'other'", - ); - assert_eq!( - modified_tool_arguments( - &json!({"choices": [{"message": {"content": "[]"}}]}), - "weather_lookup", - ) - .unwrap(), - None - ); - assert_eq!( - modified_tool_arguments( - &json!({"choices": [{"message": {"content": "{\"tool_name\":\"weather_lookup\"}"}}]}), - "weather_lookup", - ) - .unwrap(), - None - ); -} - -#[test] -fn modified_tool_payload_helpers_cover_odd_remote_payload_shapes() { - assert_flow_error_contains( - first_choice_message(&json!({"choices": [{"message": []}]})).map(|_| ()), - "did not contain choices[0].message", - ); - assert_eq!( - modified_tool_arguments( - &json!({"choices": [{"message": {"tool_calls": ["not-an-object"]}}]}), - "weather_lookup", - ) - .unwrap(), - None - ); - assert_flow_error_contains( - modified_tool_arguments( - &json!({"choices": [{"message": {"tool_calls": [{"function": {"name": "weather_lookup", "arguments": {"city": "Paris"}}}]}}]}), - "weather_lookup", - ), - "without function.arguments", - ); - assert_eq!( - modified_tool_result( - &json!({"choices": [{"message": {"role": "assistant", "content": "{\"tool_name\":\"weather_lookup\",\"result\":null}"}}]}), - "weather_lookup", - ) - .unwrap(), - Some(Json::Null) - ); - assert_flow_error_contains( - modified_tool_result( - &json!({"choices": [{"message": {"role": "tool", "name": "weather_lookup", "content": {"forecast": "rain"}}}]}), - "weather_lookup", - ), - "without message.content", - ); -} - -#[test] -fn modified_tool_result_parsing_covers_success_and_error_shapes() { - let response = json!({ - "choices": [{ - "message": { - "role": "tool", - "name": "weather_lookup", - "content": "{\"forecast\":\"cloudy\"}" - } - }] - }); - assert_eq!( - modified_tool_result(&response, "weather_lookup").unwrap(), - Some(json!({"forecast": "cloudy"})) - ); - - assert_flow_error_contains( - modified_tool_result( - &json!({"choices": [{"message": {"role": "tool", "name": "other", "content": "{}"}}]}), - "weather_lookup", - ), - "unexpected tool 'other'", - ); - assert_flow_error_contains( - modified_tool_result( - &json!({"choices": [{"message": {"role": "tool", "name": "weather_lookup"}}]}), - "weather_lookup", - ), - "without message.content", - ); - assert_flow_error_contains( - modified_tool_result( - &json!({"choices": [{"message": {"role": "tool", "name": "weather_lookup", "content": "not json"}}]}), - "weather_lookup", - ), - "not valid JSON", - ); - - let legacy = json!({ - "choices": [{ - "message": { - "content": "{\"tool_name\":\"weather_lookup\",\"result\":{\"forecast\":\"rain\"}}" - } - }] - }); - assert_eq!( - modified_tool_result(&legacy, "weather_lookup").unwrap(), - Some(json!({"forecast": "rain"})) - ); - assert_eq!( - modified_tool_result( - &json!({"choices": [{"message": {"content": "{\"tool_name\":\"weather_lookup\"}"}}]}), - "weather_lookup", - ) - .unwrap(), - None - ); - assert_eq!( - modified_tool_result( - &json!({"choices": [{"message": {"content": "[]"}}]}), - "weather_lookup", - ) - .unwrap(), - None - ); - assert_eq!( - modified_tool_result( - &json!({"choices": [{"message": {"content": "not json"}}]}), - "weather_lookup", - ) - .unwrap(), - None - ); -} - -#[test] -fn blocking_and_mark_helpers_cover_optional_payload_shapes() { - let stopped = json!({ - "guardrails": {"log": {"activated_rails": [{"name": "stop rail", "stop": true}]}} - }); - assert_eq!(blocking_rail_name(&stopped), Some("stop rail".to_string())); - - let refused = json!({ - "guardrails": { - "log": { - "activated_rails": [{ - "name": "refuse rail", - "decisions": ["refuse answer"] - }] - } - } - }); - assert_eq!( - blocking_rail_name(&refused), - Some("refuse rail".to_string()) - ); - assert_eq!( - blocking_rail_name( - &json!({"guardrails": {"log": {"activated_rails": [{"name": "allow"}]}}}) - ), - None - ); - - let mark = remote_mark_data( - true, - &Some("primary".to_string()), - &["fallback".to_string()], - Some(503), - Some("redacted".to_string()), - ); - assert_eq!(mark["stream"], Json::Bool(true)); - assert_eq!(mark["config_id"], json!("primary")); - assert_eq!(mark["config_ids"], json!(["fallback"])); - assert_eq!(mark["http_status"], json!(503)); - assert_eq!(mark["error"], json!("redacted")); - - let tool_mark = tool_remote_mark_data( - RemoteCheckKind::Output, - "weather_lookup", - &None, - &[], - Some(200), - None, - ); - assert_eq!(tool_mark["surface"], json!("tool_output")); - assert_eq!(tool_mark["tool_name"], json!("weather_lookup")); - assert_eq!( - redact_remote_error_payload(500, "sensitive body"), - "remote request failed with status 500; error body omitted from marks (14 bytes)" - ); -} - -#[tokio::test] -#[allow(clippy::await_holding_lock)] -async fn remote_execute_reports_non_stream_success_http_errors_and_invalid_json() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - crate::shared_runtime::reset_runtime_owner_for_tests(); - crate::api::runtime::set_thread_scope_stack(crate::api::runtime::create_scope_stack()); - - let success = runtime_with_endpoint(spawn_json_response(json!({ - "id": "chatcmpl-remote", - "object": "chat.completion", - "created": 1, - "model": "gpt-4o-mini", - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": "guarded"}, - "finish_reason": "stop" - }] - }))); - let response = success - .execute(simple_chat_request(), false) - .await - .expect("remote success"); - assert_eq!( - response["choices"][0]["message"]["content"], - json!("guarded") - ); - assert_eq!(success.access_state.load(Ordering::Acquire), 2); - - let invalid_json = runtime_with_endpoint(spawn_http_response( - "200 OK", - "application/json", - "not json", - )); - assert_flow_error_contains( - invalid_json.execute(simple_chat_request(), false).await, - "failed to parse remote response JSON", - ); - assert_eq!(invalid_json.access_state.load(Ordering::Acquire), 2); - - let http_error = runtime_with_endpoint(spawn_http_response( - "502 Bad Gateway", - "application/json", - r#"{"error":"backend unavailable"}"#, - )); - assert_flow_error_contains( - http_error.execute(simple_chat_request(), false).await, - "status 502", - ); - assert_eq!(http_error.access_state.load(Ordering::Acquire), 1); -} - -#[tokio::test] -#[allow(clippy::await_holding_lock)] -async fn remote_execute_transport_and_stream_status_errors_are_reported() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - crate::shared_runtime::reset_runtime_owner_for_tests(); - crate::api::runtime::set_thread_scope_stack(crate::api::runtime::create_scope_stack()); - - let runtime = RemoteBackendRuntime::new(&runtime_config(RemoteBackendConfig { - endpoint: Some(spawn_disconnecting_endpoint()), - timeout_millis: 50, - ..valid_remote() - })) - .unwrap(); - assert_flow_error_contains( - runtime.execute(simple_chat_request(), false).await, - "remote request failed", - ); - assert_eq!(runtime.access_state.load(Ordering::Acquire), 1); - assert_flow_error_contains( - runtime.execute_stream(simple_chat_request()).await, - "remote stream request failed", - ); - - let stream_status_error = runtime_with_endpoint(spawn_http_response( - "503 Service Unavailable", - "text/plain", - "downstream unavailable", - )); - assert_flow_error_contains( - stream_status_error - .execute_stream(simple_chat_request()) - .await, - "status 503", - ); -} - -#[tokio::test] -#[allow(clippy::await_holding_lock)] -async fn truncated_remote_bodies_report_buffered_streaming_and_tool_errors() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - crate::shared_runtime::reset_runtime_owner_for_tests(); - crate::api::runtime::set_thread_scope_stack(crate::api::runtime::create_scope_stack()); - - let buffered = - runtime_with_endpoint(spawn_truncated_http_response("200 OK", "application/json")); - assert_flow_error_contains( - buffered.execute(simple_chat_request(), false).await, - "failed to read remote response body", - ); - - let status = runtime_with_endpoint(spawn_truncated_http_response( - "503 Service Unavailable", - "text/plain", - )); - assert_flow_error_contains( - status.execute_stream(simple_chat_request()).await, - "failed to read remote stream error body", - ); - - let streaming = - runtime_with_endpoint(spawn_truncated_http_response("200 OK", "text/event-stream")); - let mut stream = streaming - .execute_stream(simple_chat_request()) - .await - .expect("stream opens after headers"); - assert_flow_error_contains( - stream - .next() - .await - .expect("truncated body reports an error"), - "failed to read remote stream chunk", - ); - - let tool = runtime_with_endpoint(spawn_truncated_http_response("200 OK", "application/json")); - assert_flow_error_contains( - tool.check_tool_input("weather_lookup", &json!({"city": "Phoenix"})) - .await, - "failed to read remote response body", - ); -} - -#[tokio::test] -#[allow(clippy::await_holding_lock)] -async fn remote_execute_stream_yields_completed_events_and_reports_malformed_final_event() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - crate::shared_runtime::reset_runtime_owner_for_tests(); - crate::api::runtime::set_thread_scope_stack(crate::api::runtime::create_scope_stack()); - - let runtime = runtime_with_endpoint(spawn_http_response( - "200 OK", - "text/event-stream", - concat!( - "data: {\"chunk\":\"first\"}\n\n", - "event: done\ndata: {\"chunk\":\"final\"}\n\n" - ), - )); - let mut stream = runtime - .execute_stream(simple_chat_request()) - .await - .expect("remote stream"); - assert_eq!( - stream.next().await.unwrap().unwrap()["chunk"], - json!("first") - ); - assert_eq!( - stream.next().await.unwrap().unwrap()["chunk"], - json!("final") - ); - assert!(stream.next().await.is_none()); - - let malformed = runtime_with_endpoint(spawn_http_response( - "200 OK", - "text/event-stream", - "data: {\"chunk\":", - )); - let mut stream = malformed - .execute_stream(simple_chat_request()) - .await - .expect("malformed stream opens"); - assert_flow_error_contains( - stream - .next() - .await - .expect("decoder should report final frame"), - "failed to parse SSE data payload", - ); -} - -#[tokio::test] -#[allow(clippy::await_holding_lock)] -async fn remote_execute_stream_reports_malformed_named_final_event() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - crate::shared_runtime::reset_runtime_owner_for_tests(); - crate::api::runtime::set_thread_scope_stack(crate::api::runtime::create_scope_stack()); - - let runtime = runtime_with_endpoint(spawn_http_response( - "200 OK", - "text/event-stream", - concat!(": keep-alive\n\n", "event: done\ndata: not-json\n\n"), - )); - let mut stream = runtime - .execute_stream(simple_chat_request()) - .await - .expect("remote stream"); - assert_flow_error_contains( - stream - .next() - .await - .expect("decoder should report malformed done event"), - "failed to parse SSE data payload", - ); -} - -#[tokio::test] -#[allow(clippy::await_holding_lock)] -async fn remote_stream_decoder_flushes_an_unterminated_final_event() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - crate::shared_runtime::reset_runtime_owner_for_tests(); - crate::api::runtime::set_thread_scope_stack(crate::api::runtime::create_scope_stack()); - - let runtime = runtime_with_endpoint(spawn_http_response( - "200 OK", - "text/event-stream", - "data: {\"final\":true}", - )); - let mut stream = runtime.execute_stream(simple_chat_request()).await.unwrap(); - assert_eq!(stream.next().await.unwrap().unwrap()["final"], json!(true)); - assert!(stream.next().await.is_none()); -} - -#[tokio::test] -#[allow(clippy::await_holding_lock)] -async fn remote_tool_input_checks_cover_rewrite_block_noop_and_invalid_json() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - crate::shared_runtime::reset_runtime_owner_for_tests(); - crate::api::runtime::set_thread_scope_stack(crate::api::runtime::create_scope_stack()); - - let rewritten = runtime_with_endpoint(spawn_json_response(json!({ - "choices": [{ - "message": { - "tool_calls": [{ - "function": { - "name": "weather_lookup", - "arguments": "{\"city\":\"Paris\"}" - } - }] - } - }] - }))) - .check_tool_input("weather_lookup", &json!({"city": "Phoenix"})) - .await - .expect("modified tool input"); - assert_eq!(rewritten, json!({"city": "Paris"})); - - let blocked = runtime_with_endpoint(spawn_json_response(json!({ - "guardrails": { - "log": { - "activated_rails": [{"name": "input rail", "stop": true}] - } - } - }))); - assert_flow_error_contains( - blocked - .check_tool_input("weather_lookup", &json!({"city": "Phoenix"})) - .await, - "tool_input rail blocked", - ); - - let original = json!({"city": "Phoenix"}); - let noop = runtime_with_endpoint(spawn_json_response(json!({ - "choices": [{"message": {"role": "assistant", "content": ""}}], - "guardrails": {"log": {"activated_rails": []}} - }))) - .check_tool_input("weather_lookup", &original) - .await - .expect("noop tool input"); - assert_eq!(noop, original); - - let invalid_json = runtime_with_endpoint(spawn_http_response( - "200 OK", - "application/json", - "not json", - )); - assert_flow_error_contains( - invalid_json - .check_tool_input("weather_lookup", &json!({"city": "Phoenix"})) - .await, - "failed to parse remote response JSON", - ); -} - -#[tokio::test] -#[allow(clippy::await_holding_lock)] -async fn remote_tool_output_checks_cover_rewrite_block_and_noop() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - crate::shared_runtime::reset_runtime_owner_for_tests(); - crate::api::runtime::set_thread_scope_stack(crate::api::runtime::create_scope_stack()); - - let rewritten = runtime_with_endpoint(spawn_json_response(json!({ - "choices": [{ - "message": { - "role": "tool", - "name": "weather_lookup", - "content": "{\"forecast\":\"rain\"}" - } - }] - }))) - .check_tool_output( - "weather_lookup", - &json!({"city": "Phoenix"}), - &json!({"forecast": "sunny"}), - ) - .await - .expect("modified tool output"); - assert_eq!(rewritten, json!({"forecast": "rain"})); - - let blocked = runtime_with_endpoint(spawn_json_response(json!({ - "guardrails": { - "log": { - "activated_rails": [{ - "name": "output rail", - "decisions": ["execute check", "refuse answer"] - }] - } - } - }))); - assert_flow_error_contains( - blocked - .check_tool_output( - "weather_lookup", - &json!({"city": "Phoenix"}), - &json!({"forecast": "sunny"}), - ) - .await, - "tool_output rail blocked", - ); - - let original = json!({"forecast": "sunny"}); - let noop = runtime_with_endpoint(spawn_json_response(json!({ - "choices": [{"message": {"role": "assistant", "content": ""}}], - "guardrails": {"log": {"activated_rails": []}} - }))) - .check_tool_output("weather_lookup", &json!({"city": "Phoenix"}), &original) - .await - .expect("noop tool output"); - assert_eq!(noop, original); -} - -#[tokio::test] -#[allow(clippy::await_holding_lock)] -async fn tool_remote_check_transport_failures_are_reported() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - crate::shared_runtime::reset_runtime_owner_for_tests(); - let stack = crate::api::runtime::create_scope_stack(); - crate::api::runtime::set_thread_scope_stack(stack); - - let runtime = RemoteBackendRuntime::new(&runtime_config(RemoteBackendConfig { - endpoint: Some(spawn_disconnecting_endpoint()), - timeout_millis: 50, - ..valid_remote() - })) - .unwrap(); - assert_flow_error_contains( - runtime - .check_tool_input("weather_lookup", &json!({"city": "Phoenix"})) - .await, - "remote request failed", - ); - assert_flow_error_contains( - runtime - .check_tool_output( - "weather_lookup", - &json!({"city": "Phoenix"}), - &json!({"forecast": "sunny"}), - ) - .await, - "remote request failed", - ); -} - -#[tokio::test] -#[allow(clippy::await_holding_lock)] -async fn tool_remote_check_http_status_failures_are_reported() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - crate::shared_runtime::reset_runtime_owner_for_tests(); - crate::api::runtime::set_thread_scope_stack(crate::api::runtime::create_scope_stack()); - - let input_status_error = runtime_with_endpoint(spawn_http_response( - "429 Too Many Requests", - "application/json", - r#"{"error":"limited"}"#, - )); - assert_flow_error_contains( - input_status_error - .check_tool_input("weather_lookup", &json!({"city": "Phoenix"})) - .await, - "status 429", - ); - - let output_status_error = runtime_with_endpoint(spawn_http_response( - "503 Service Unavailable", - "application/json", - r#"{"error":"maintenance"}"#, - )); - assert_flow_error_contains( - output_status_error - .check_tool_output( - "weather_lookup", - &json!({"city": "Phoenix"}), - &json!({"forecast": "sunny"}), - ) - .await, - "status 503", - ); - - let permission_error = runtime_with_endpoint(spawn_http_response( - "403 Forbidden", - "application/json", - r#"{"error":"forbidden"}"#, - )); - assert_flow_error_contains( - permission_error - .check_tool_input("weather_lookup", &json!({"city": "Phoenix"})) - .await, - "status 403", - ); -} - -#[tokio::test] -#[allow(clippy::await_holding_lock)] -async fn registered_remote_tool_input_intercept_rewrites_args_and_skips_output_check() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - crate::shared_runtime::reset_runtime_owner_for_tests(); - crate::api::runtime::set_thread_scope_stack(crate::api::runtime::create_scope_stack()); - - let endpoint = spawn_json_response(json!({ - "choices": [{ - "message": { - "tool_calls": [{ - "function": { - "name": "weather_lookup", - "arguments": "{\"city\":\"Paris\"}" - } - }] - } - }], - "guardrails": {"log": {"activated_rails": []}} - })); - let mut ctx = crate::plugin::PluginRegistrationContext::new(); - register_remote_backend( - NeMoGuardrailsConfig { - input: false, - output: false, - tool_input: true, - tool_output: false, - remote: Some(RemoteBackendConfig { - endpoint: Some(endpoint), - config_id: Some("safety-default".to_string()), - timeout_millis: 5_000, - ..RemoteBackendConfig::default() - }), - ..NeMoGuardrailsConfig::default() - }, - &mut ctx, - ) - .unwrap(); - let mut registrations = ctx.into_registrations(); - - let callback_args = Arc::new(std::sync::Mutex::new(Json::Null)); - let seen = Arc::clone(&callback_args); - let callback: crate::api::runtime::ToolExecutionNextFn = Arc::new(move |args| { - let seen = Arc::clone(&seen); - Box::pin(async move { - *seen.lock().unwrap() = args; - Ok(json!({"forecast": "sunny"}).into()) - }) - }); - - let result = crate::api::tool::tool_call_execute( - crate::api::tool::ToolCallExecuteParams::builder() - .name("weather_lookup") - .args(json!({"city": "Phoenix"})) - .func(callback) - .build(), - ) - .await - .unwrap(); - - assert_eq!(*callback_args.lock().unwrap(), json!({"city": "Paris"})); - assert_eq!(result.result, json!({"forecast": "sunny"})); - - crate::plugin::rollback_registrations(&mut registrations); -} diff --git a/crates/core/tests/unit/plugins/nemo_guardrails/remote_tests.rs b/crates/core/tests/unit/plugins/nemo_guardrails/remote_tests.rs deleted file mode 100644 index c3025bdf61..0000000000 --- a/crates/core/tests/unit/plugins/nemo_guardrails/remote_tests.rs +++ /dev/null @@ -1,1500 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Remote runtime tests for the NeMo Guardrails plugin component. -#![allow(clippy::await_holding_lock)] - -use super::*; - -#[tokio::test] -async fn remote_initialization_installs_non_streaming_execution_intercept() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - let events = capture_events("nemo-guardrails-remote-execution-events"); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, request_rx) = mpsc::channel(); - let response_body = json!({ - "id": "chatcmpl-remote", - "object": "chat.completion", - "created": 1, - "model": "gpt-4o-mini", - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": "guarded"}, - "finish_reason": "stop" - }], - "guardrails": { - "config_id": "safety-default", - "state": {"state": {"conversation": "server-state"}}, - "output_data": {"decision": "allow"} - } - }) - .to_string(); - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default", - "headers": {"x-guardrails-auth": "token"}, - "timeout_millis": 5_000 - }, - "request_defaults": { - "context": {"tenant": "acme"}, - "thread_id": "thread-1234567890", - "state": {"state": {"conversation": "client-state"}}, - "rails": {"input": true, "retrieval": ["kb"]}, - "llm_params": {"temperature": 0.1}, - "llm_output": true, - "output_vars": ["answer"], - "log": {"activated_rails": true} - } - }))) - .await - .unwrap(); - - let original_called = Arc::new(AtomicBool::new(false)); - let called = Arc::clone(&original_called); - let func: LlmExecutionNextFn = Arc::new(move |_req| { - called.store(true, Ordering::SeqCst); - Box::pin(async move { Ok(json!({"response": "original"})) }) - }); - - let response = llm_call_execute( - LlmCallExecuteParams::builder() - .name("openai") - .request(make_chat_request(false)) - .func(func) - .attributes(LlmAttributes::empty()) - .response_codec(Arc::new(OpenAIChatCodec) as Arc) - .build(), - ) - .await - .unwrap(); - - assert!(!original_called.load(Ordering::SeqCst)); - assert_eq!(response["id"], json!("chatcmpl-remote")); - assert_eq!(response["object"], json!("chat.completion")); - assert_eq!(response["model"], json!("gpt-4o-mini")); - assert_eq!( - response["choices"][0]["message"]["content"], - json!("guarded") - ); - assert_eq!( - response["guardrails"]["output_data"]["decision"], - json!("allow") - ); - assert_eq!( - response["guardrails"]["state"]["state"]["conversation"], - json!("server-state") - ); - - let captured = recv_captured_request(&request_rx); - assert_eq!(captured.path, "/v1/chat/completions"); - assert!(captured.content_type.starts_with("application/json")); - - let request_json: Json = serde_json::from_slice(&captured.body).unwrap(); - assert_eq!(request_json["messages"][0]["content"], json!("hello")); - assert_eq!(request_json["stream"], json!(false)); - assert_eq!( - request_json["guardrails"]["config_id"], - json!("safety-default") - ); - assert_eq!( - request_json["guardrails"]["context"]["tenant"], - json!("acme") - ); - assert_eq!( - request_json["guardrails"]["thread_id"], - json!("thread-1234567890") - ); - assert_eq!( - request_json["guardrails"]["state"]["state"]["conversation"], - json!("client-state") - ); - assert_eq!( - request_json["guardrails"]["options"]["rails"]["retrieval"], - json!(["kb"]) - ); - assert_eq!( - request_json["guardrails"]["options"]["llm_output"], - json!(true) - ); - - let captured_events = captured_events_snapshot(&events); - let mark_names: Vec<_> = captured_events - .iter() - .filter(|event| event.kind() == "mark") - .map(|event| event.name().to_string()) - .collect(); - assert!(mark_names.contains(&"nemo_guardrails.remote.start".to_string())); - assert!(mark_names.contains(&"nemo_guardrails.remote.end".to_string())); - - let start_mark = captured_events - .iter() - .find(|event| event.name() == "nemo_guardrails.remote.start") - .unwrap(); - assert_eq!( - start_mark.data().unwrap()["config_id"], - json!("safety-default") - ); - assert_eq!(start_mark.data().unwrap()["stream"], json!(false)); - - let end_mark = captured_events - .iter() - .find(|event| event.name() == "nemo_guardrails.remote.end") - .unwrap(); - assert_eq!(end_mark.data().unwrap()["http_status"], json!(200)); - assert_eq!(end_mark.data().unwrap()["stream"], json!(false)); - - deregister_subscriber("nemo-guardrails-remote-execution-events").unwrap(); -} - -#[tokio::test] -async fn remote_request_uses_config_ids_when_config_id_is_not_set() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, request_rx) = mpsc::channel(); - let response_body = json!({ - "id": "chatcmpl-remote", - "object": "chat.completion", - "created": 1, - "model": "gpt-4o-mini", - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": "guarded"}, - "finish_reason": "stop" - }] - }) - .to_string(); - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": format!("http://{address}"), - "config_ids": ["safety-a", "safety-b"] - } - }))) - .await - .unwrap(); - - let func: LlmExecutionNextFn = - Arc::new(move |_req| Box::pin(async move { Ok(json!({"response": "original"})) })); - - let _ = llm_call_execute( - LlmCallExecuteParams::builder() - .name("openai") - .request(make_chat_request(false)) - .func(func) - .attributes(LlmAttributes::empty()) - .response_codec(Arc::new(OpenAIChatCodec) as Arc) - .build(), - ) - .await - .unwrap(); - - let captured = recv_captured_request(&request_rx); - let request_json: Json = serde_json::from_slice(&captured.body).unwrap(); - assert_eq!( - request_json["guardrails"]["config_ids"], - json!(["safety-a", "safety-b"]) - ); - assert!(request_json["guardrails"].get("config_id").is_none()); -} - -#[tokio::test] -async fn remote_llm_request_disables_input_rails_when_surface_is_off() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, request_rx) = mpsc::channel(); - let response_body = json!({ - "id": "chatcmpl-remote", - "object": "chat.completion", - "created": 1, - "model": "gpt-4o-mini", - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": "guarded"}, - "finish_reason": "stop" - }] - }) - .to_string(); - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "input": false, - "output": true, - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - }, - "request_defaults": { - "rails": { - "input": ["self check input"], - "output": ["self check output"], - "retrieval": ["kb"] - } - } - }))) - .await - .unwrap(); - - let func: LlmExecutionNextFn = - Arc::new(move |_req| Box::pin(async move { Ok(json!({"response": "original"})) })); - - let _ = llm_call_execute( - LlmCallExecuteParams::builder() - .name("openai") - .request(make_chat_request(false)) - .func(func) - .attributes(LlmAttributes::empty()) - .response_codec(Arc::new(OpenAIChatCodec) as Arc) - .build(), - ) - .await - .unwrap(); - - let captured = recv_captured_request(&request_rx); - let request_json: Json = serde_json::from_slice(&captured.body).unwrap(); - assert_eq!( - request_json["guardrails"]["options"]["rails"]["input"], - json!(false) - ); - assert_eq!( - request_json["guardrails"]["options"]["rails"]["output"], - json!(["self check output"]) - ); - assert_eq!( - request_json["guardrails"]["options"]["rails"]["retrieval"], - json!(["kb"]) - ); -} - -#[tokio::test] -async fn remote_llm_request_disables_output_rails_when_surface_is_off() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, request_rx) = mpsc::channel(); - let response_body = json!({ - "id": "chatcmpl-remote", - "object": "chat.completion", - "created": 1, - "model": "gpt-4o-mini", - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": "guarded"}, - "finish_reason": "stop" - }] - }) - .to_string(); - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "input": true, - "output": false, - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - }, - "request_defaults": { - "rails": { - "input": ["self check input"], - "output": ["self check output"], - "dialog": true - } - } - }))) - .await - .unwrap(); - - let func: LlmExecutionNextFn = - Arc::new(move |_req| Box::pin(async move { Ok(json!({"response": "original"})) })); - - let _ = llm_call_execute( - LlmCallExecuteParams::builder() - .name("openai") - .request(make_chat_request(false)) - .func(func) - .attributes(LlmAttributes::empty()) - .response_codec(Arc::new(OpenAIChatCodec) as Arc) - .build(), - ) - .await - .unwrap(); - - let captured = recv_captured_request(&request_rx); - let request_json: Json = serde_json::from_slice(&captured.body).unwrap(); - assert_eq!( - request_json["guardrails"]["options"]["rails"]["input"], - json!(["self check input"]) - ); - assert_eq!( - request_json["guardrails"]["options"]["rails"]["output"], - json!(false) - ); - assert_eq!( - request_json["guardrails"]["options"]["rails"]["dialog"], - json!(true) - ); -} - -#[tokio::test] -async fn remote_initialization_installs_stream_execution_intercept() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - let events = capture_events("nemo-guardrails-remote-stream-events"); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, request_rx) = mpsc::channel(); - let sse_body = concat!( - "data: {\"id\":\"chatcmpl-remote\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"gpt-4o-mini\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"guard\"},\"finish_reason\":null}]}\n\n", - "data: {\"id\":\"chatcmpl-remote\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"gpt-4o-mini\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ed\"},\"finish_reason\":\"stop\"}]}\n\n", - "data: [DONE]\n\n" - ); - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\n\r\n{}", - sse_body.len(), - sse_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - } - }))) - .await - .unwrap(); - - let original_called = Arc::new(AtomicBool::new(false)); - let called = Arc::clone(&original_called); - let func: LlmStreamExecutionNextFn = Arc::new(move |_req| { - called.store(true, Ordering::SeqCst); - Box::pin(async move { - let stream = tokio_stream::iter(vec![Ok(json!({"chunk": "original"}))]); - Ok(LlmJsonStream::new(stream)) - }) - }); - - let streaming_codec = OpenAIChatStreamingCodec::new(); - let collector = streaming_codec.collector(); - let finalizer = streaming_codec.finalizer(); - let response_codec: Arc = Arc::new(OpenAIChatCodec); - - let mut stream = llm_stream_call_execute( - LlmStreamCallExecuteParams::builder() - .name("openai") - .request(make_chat_request(true)) - .func(func) - .collector(collector) - .finalizer(finalizer) - .attributes(LlmAttributes::STREAMING) - .response_codec(response_codec) - .build(), - ) - .await - .unwrap(); - - let mut chunks = Vec::new(); - while let Some(chunk) = tokio::time::timeout(TEST_TIMEOUT, stream.next()) - .await - .expect("timed out waiting for remote stream chunk") - { - chunks.push(chunk.unwrap()); - } - - assert!(!original_called.load(Ordering::SeqCst)); - assert_eq!(chunks.len(), 2); - assert_eq!(chunks[0]["choices"][0]["delta"]["content"], json!("guard")); - assert_eq!(chunks[1]["choices"][0]["delta"]["content"], json!("ed")); - - let captured = recv_captured_request(&request_rx); - let request_json: Json = serde_json::from_slice(&captured.body).unwrap(); - assert_eq!(request_json["stream"], json!(true)); - assert_eq!( - request_json["guardrails"]["config_id"], - json!("safety-default") - ); - - let captured_events = captured_events_snapshot(&events); - let start_mark = captured_events - .iter() - .find(|event| event.name() == "nemo_guardrails.remote.start") - .unwrap(); - assert_eq!(start_mark.data().unwrap()["stream"], json!(true)); - - let end_mark = captured_events - .iter() - .find(|event| event.name() == "nemo_guardrails.remote.end") - .unwrap(); - assert_eq!(end_mark.data().unwrap()["http_status"], json!(200)); - assert_eq!(end_mark.data().unwrap()["stream"], json!(true)); - - deregister_subscriber("nemo-guardrails-remote-stream-events").unwrap(); -} - -#[tokio::test] -async fn remote_non_streaming_http_errors_are_reported_and_marked() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - let events = capture_events("nemo-guardrails-remote-error-events"); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, _request_rx) = mpsc::channel(); - let response_body = r#"{"error":"backend unavailable"}"#; - let http_response = format!( - "HTTP/1.1 502 Bad Gateway\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - } - }))) - .await - .unwrap(); - - let original_called = Arc::new(AtomicBool::new(false)); - let called = Arc::clone(&original_called); - let func: LlmExecutionNextFn = Arc::new(move |_req| { - called.store(true, Ordering::SeqCst); - Box::pin(async move { Ok(json!({"response": "original"})) }) - }); - - let error = llm_call_execute( - LlmCallExecuteParams::builder() - .name("openai") - .request(make_chat_request(false)) - .func(func) - .attributes(LlmAttributes::empty()) - .response_codec(Arc::new(OpenAIChatCodec) as Arc) - .build(), - ) - .await - .unwrap_err(); - - assert!(!original_called.load(Ordering::SeqCst)); - match error { - crate::error::FlowError::Internal(message) => { - assert!(message.contains("status 502")); - assert!(message.contains("backend unavailable")); - } - other => panic!("unexpected error: {other}"), - } - - let captured_events = captured_events_snapshot(&events); - assert!( - captured_events - .iter() - .any(|event| event.name() == "nemo_guardrails.remote.start") - ); - let error_mark = captured_events - .iter() - .find(|event| event.name() == "nemo_guardrails.remote.error") - .unwrap(); - assert_eq!(error_mark.data().unwrap()["http_status"], json!(502)); - assert_eq!(error_mark.data().unwrap()["stream"], json!(false)); - assert!( - error_mark.data().unwrap()["error"] - .as_str() - .unwrap() - .contains("error body omitted from marks") - ); - - deregister_subscriber("nemo-guardrails-remote-error-events").unwrap(); -} - -#[tokio::test] -async fn remote_streaming_http_errors_are_reported_and_marked() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - let events = capture_events("nemo-guardrails-remote-stream-error-events"); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, _request_rx) = mpsc::channel(); - let response_body = r#"{"error":"stream backend unavailable"}"#; - let http_response = format!( - "HTTP/1.1 503 Service Unavailable\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - } - }))) - .await - .unwrap(); - - let original_called = Arc::new(AtomicBool::new(false)); - let called = Arc::clone(&original_called); - let func: LlmStreamExecutionNextFn = Arc::new(move |_req| { - called.store(true, Ordering::SeqCst); - Box::pin(async move { - let stream = tokio_stream::iter(vec![Ok(json!({"chunk": "original"}))]); - Ok(LlmJsonStream::new(stream)) - }) - }); - - let streaming_codec = OpenAIChatStreamingCodec::new(); - let collector = streaming_codec.collector(); - let finalizer = streaming_codec.finalizer(); - let response_codec: Arc = Arc::new(OpenAIChatCodec); - - let error = match llm_stream_call_execute( - LlmStreamCallExecuteParams::builder() - .name("openai") - .request(make_chat_request(true)) - .func(func) - .collector(collector) - .finalizer(finalizer) - .attributes(LlmAttributes::STREAMING) - .response_codec(response_codec) - .build(), - ) - .await - { - Ok(_) => panic!("expected remote streaming request to fail"), - Err(error) => error, - }; - - assert!(!original_called.load(Ordering::SeqCst)); - match error { - crate::error::FlowError::Internal(message) => { - assert!(message.contains("status 503")); - assert!(message.contains("stream backend unavailable")); - } - other => panic!("unexpected error: {other}"), - } - - let captured_events = captured_events_snapshot(&events); - assert!( - captured_events - .iter() - .any(|event| event.name() == "nemo_guardrails.remote.start") - ); - let error_mark = captured_events - .iter() - .find(|event| event.name() == "nemo_guardrails.remote.error") - .unwrap(); - assert_eq!(error_mark.data().unwrap()["http_status"], json!(503)); - assert_eq!(error_mark.data().unwrap()["stream"], json!(true)); - assert!( - error_mark.data().unwrap()["error"] - .as_str() - .unwrap() - .contains("error body omitted from marks") - ); - - deregister_subscriber("nemo-guardrails-remote-stream-error-events").unwrap(); -} - -#[tokio::test] -async fn remote_non_streaming_invalid_json_is_reported_and_marked() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - let events = capture_events("nemo-guardrails-remote-invalid-json-events"); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, _request_rx) = mpsc::channel(); - let response_body = "{not-json}"; - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - } - }))) - .await - .unwrap(); - - let func: LlmExecutionNextFn = - Arc::new(move |_req| Box::pin(async move { Ok(json!({"response": "original"})) })); - - let error = llm_call_execute( - LlmCallExecuteParams::builder() - .name("openai") - .request(make_chat_request(false)) - .func(func) - .attributes(LlmAttributes::empty()) - .response_codec(Arc::new(OpenAIChatCodec) as Arc) - .build(), - ) - .await - .unwrap_err(); - - match error { - crate::error::FlowError::Internal(message) => { - assert!(message.contains("failed to parse remote response JSON")); - } - other => panic!("unexpected error: {other}"), - } - - let captured_events = captured_events_snapshot(&events); - let error_mark = captured_events - .iter() - .find(|event| event.name() == "nemo_guardrails.remote.error") - .unwrap(); - assert_eq!(error_mark.data().unwrap()["http_status"], json!(200)); - assert_eq!(error_mark.data().unwrap()["stream"], json!(false)); - - deregister_subscriber("nemo-guardrails-remote-invalid-json-events").unwrap(); -} - -#[tokio::test] -async fn remote_streaming_malformed_chunk_is_reported_and_marked() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - let events = capture_events("nemo-guardrails-remote-malformed-stream-events"); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, _request_rx) = mpsc::channel(); - let sse_body = "data: {not-json}\n\n"; - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\n\r\n{}", - sse_body.len(), - sse_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - } - }))) - .await - .unwrap(); - - let func: LlmStreamExecutionNextFn = Arc::new(move |_req| { - Box::pin(async move { - let stream = tokio_stream::iter(vec![Ok(json!({"chunk": "original"}))]); - Ok(LlmJsonStream::new(stream)) - }) - }); - - let streaming_codec = OpenAIChatStreamingCodec::new(); - let collector = streaming_codec.collector(); - let finalizer = streaming_codec.finalizer(); - let response_codec: Arc = Arc::new(OpenAIChatCodec); - - let mut stream = llm_stream_call_execute( - LlmStreamCallExecuteParams::builder() - .name("openai") - .request(make_chat_request(true)) - .func(func) - .collector(collector) - .finalizer(finalizer) - .attributes(LlmAttributes::STREAMING) - .response_codec(response_codec) - .build(), - ) - .await - .unwrap(); - - let error = tokio::time::timeout(TEST_TIMEOUT, stream.next()) - .await - .expect("timed out waiting for remote stream error") - .unwrap() - .unwrap_err(); - match error { - crate::error::FlowError::Internal(message) => { - assert!(!message.is_empty()); - } - other => panic!("unexpected error: {other}"), - } - - let captured_events = captured_events_snapshot(&events); - let error_mark = captured_events - .iter() - .find(|event| event.name() == "nemo_guardrails.remote.error") - .unwrap(); - assert_eq!(error_mark.data().unwrap()["http_status"], json!(200)); - assert_eq!(error_mark.data().unwrap()["stream"], json!(true)); - - deregister_subscriber("nemo-guardrails-remote-malformed-stream-events").unwrap(); -} - -#[tokio::test] -async fn remote_preflight_tool_choice_failure_is_reported_and_marked() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - let events = capture_events("nemo-guardrails-remote-preflight-error-events"); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": unused_local_endpoint(), - "config_id": "safety-default" - } - }))) - .await - .unwrap(); - - let func: LlmExecutionNextFn = - Arc::new(move |_req| Box::pin(async move { Ok(json!({"response": "original"})) })); - let request = LlmRequest { - headers: serde_json::Map::new(), - content: json!({ - "model": "gpt-4o-mini", - "messages": [{"role": "user", "content": "hello"}], - "tools": [{ - "type": "function", - "function": { - "name": "lookup", - "description": "Lookup data", - "parameters": {"type": "object"} - } - }] - }), - }; - - let error = llm_call_execute( - LlmCallExecuteParams::builder() - .name("openai") - .request(request) - .func(func) - .attributes(LlmAttributes::empty()) - .response_codec(Arc::new(OpenAIChatCodec) as Arc) - .build(), - ) - .await - .unwrap_err(); - - match error { - crate::error::FlowError::Internal(message) => { - assert!(message.contains("does not support OpenAI tool definitions or tool_choice")); - } - other => panic!("unexpected error: {other}"), - } - - let captured_events = captured_events_snapshot(&events); - assert!( - captured_events - .iter() - .any(|event| event.name() == "nemo_guardrails.remote.start") - ); - let error_mark = captured_events - .iter() - .find(|event| event.name() == "nemo_guardrails.remote.error") - .unwrap(); - assert_eq!(error_mark.data().unwrap()["stream"], json!(false)); - assert!( - error_mark.data().unwrap()["error"] - .as_str() - .unwrap() - .contains("does not support OpenAI tool definitions or tool_choice") - ); - - deregister_subscriber("nemo-guardrails-remote-preflight-error-events").unwrap(); -} - -#[tokio::test] -async fn remote_transport_failure_is_reported_and_marked() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - let events = capture_events("nemo-guardrails-remote-transport-error-events"); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": unused_local_endpoint(), - "config_id": "safety-default", - "timeout_millis": 50 - } - }))) - .await - .unwrap(); - - let func: LlmExecutionNextFn = - Arc::new(move |_req| Box::pin(async move { Ok(json!({"response": "original"})) })); - - let error = llm_call_execute( - LlmCallExecuteParams::builder() - .name("openai") - .request(make_chat_request(false)) - .func(func) - .attributes(LlmAttributes::empty()) - .response_codec(Arc::new(OpenAIChatCodec) as Arc) - .build(), - ) - .await - .unwrap_err(); - - match error { - crate::error::FlowError::Internal(message) => { - assert!(message.contains("remote request failed")); - } - other => panic!("unexpected error: {other}"), - } - - let captured_events = captured_events_snapshot(&events); - let error_mark = captured_events - .iter() - .find(|event| event.name() == "nemo_guardrails.remote.error") - .unwrap(); - assert_eq!(error_mark.data().unwrap()["stream"], json!(false)); - assert!(error_mark.data().unwrap().get("http_status").is_none()); - - deregister_subscriber("nemo-guardrails-remote-transport-error-events").unwrap(); -} - -#[tokio::test] -async fn remote_success_without_guardrails_payload_is_allowed() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, _request_rx) = mpsc::channel(); - let response_body = json!({ - "id": "chatcmpl-remote", - "object": "chat.completion", - "created": 1, - "model": "gpt-4o-mini", - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": "guarded"}, - "finish_reason": "stop" - }] - }) - .to_string(); - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "codec": "openai_chat", - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - } - }))) - .await - .unwrap(); - - let func: LlmExecutionNextFn = - Arc::new(move |_req| Box::pin(async move { Ok(json!({"response": "original"})) })); - - let response = llm_call_execute( - LlmCallExecuteParams::builder() - .name("openai") - .request(make_chat_request(false)) - .func(func) - .attributes(LlmAttributes::empty()) - .response_codec(Arc::new(OpenAIChatCodec) as Arc) - .build(), - ) - .await - .unwrap(); - - assert_eq!(response["id"], json!("chatcmpl-remote")); - assert!(response.get("guardrails").is_none()); -} - -#[tokio::test] -async fn remote_tool_output_can_rewrite_tool_result() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, request_rx) = mpsc::channel(); - let response_body = json!({ - "id": "chatcmpl-tool-output-modified", - "object": "chat.completion", - "created": 1, - "model": "", - "choices": [{ - "index": 0, - "message": { - "role": "tool", - "name": "weather_lookup", - "tool_call_id": "call_weather_lookup_1", - "content": "{\"forecast\":\"cloudy\"}" - }, - "finish_reason": "stop" - }], - "guardrails": { - "config_id": "safety-default", - "log": { - "activated_rails": [] - } - } - }) - .to_string(); - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "input": false, - "output": false, - "tool_output": true, - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - } - }))) - .await - .unwrap(); - - let result = tool_call_execute( - ToolCallExecuteParams::builder() - .name("weather_lookup") - .args(json!({"city": "Phoenix"})) - .func(Arc::new(move |_args| { - Box::pin(async move { Ok(json!({"forecast": "sunny"}).into()) }) - })) - .build(), - ) - .await - .unwrap(); - - assert_eq!(result.result, json!({"forecast": "cloudy"})); - - let captured = recv_captured_request(&request_rx); - let request_json: Json = serde_json::from_slice(&captured.body).unwrap(); - assert_eq!( - request_json["guardrails"]["options"]["rails"]["tool_input"], - json!(true) - ); - assert_eq!( - request_json["guardrails"]["options"]["rails"]["tool_output"], - json!(false) - ); - assert_eq!(request_json["messages"][0]["role"], json!("user")); - assert_eq!(request_json["messages"][1]["role"], json!("assistant")); - assert_eq!(request_json["messages"][2]["role"], json!("tool")); -} - -#[tokio::test] -async fn remote_tool_output_rejects_when_remote_rail_refuses_without_stop_flag() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, _request_rx) = mpsc::channel(); - let response_body = json!({ - "id": "chatcmpl-tool-output-blocked", - "object": "chat.completion", - "created": 1, - "model": "", - "choices": [{ - "index": 0, - "message": { - "role": "assistant", - "content": "Tool result blocked by validation." - }, - "finish_reason": "stop" - }], - "guardrails": { - "config_id": "safety-default", - "log": { - "activated_rails": [{ - "type": "dialog", - "name": "validate tool result content", - "decisions": ["execute validate_tool_result_content", "refuse tool input"], - "stop": false - }] - } - } - }) - .to_string(); - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "input": false, - "output": false, - "tool_output": true, - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - } - }))) - .await - .unwrap(); - - let error = tool_call_execute( - ToolCallExecuteParams::builder() - .name("weather_lookup") - .args(json!({"city": "Phoenix"})) - .func(Arc::new(move |_args| { - Box::pin(async move { Ok(json!({"secret": "ZXCVSECRET42"}).into()) }) - })) - .build(), - ) - .await - .unwrap_err(); - - assert!( - error - .to_string() - .contains("tool_output rail blocked tool call by rail 'validate tool result content'") - ); -} - -#[tokio::test] -async fn remote_tool_output_preserves_named_rail_selectors() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, request_rx) = mpsc::channel(); - let response_body = json!({ - "id": "chatcmpl-tool-output-modified", - "object": "chat.completion", - "created": 1, - "model": "", - "choices": [{ - "index": 0, - "message": { - "role": "tool", - "name": "weather_lookup", - "tool_call_id": "call_weather_lookup_1", - "content": "{\"forecast\":\"cloudy\"}" - }, - "finish_reason": "stop" - }], - "guardrails": { - "config_id": "safety-default", - "log": { - "activated_rails": [] - } - } - }) - .to_string(); - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "input": false, - "output": false, - "tool_output": true, - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - }, - "request_defaults": { - "rails": { - "tool_output": ["validate_tool_output"] - } - } - }))) - .await - .unwrap(); - - let _ = tool_call_execute( - ToolCallExecuteParams::builder() - .name("weather_lookup") - .args(json!({"city": "Phoenix"})) - .func(Arc::new(move |_args| { - Box::pin(async move { Ok(json!({"forecast": "sunny"}).into()) }) - })) - .build(), - ) - .await - .unwrap(); - - let captured = recv_captured_request(&request_rx); - let request_json: Json = serde_json::from_slice(&captured.body).unwrap(); - assert_eq!( - request_json["guardrails"]["options"]["rails"]["tool_input"], - json!(["validate_tool_output"]) - ); - assert_eq!( - request_json["guardrails"]["options"]["rails"]["tool_output"], - json!(false) - ); -} - -#[tokio::test] -async fn remote_tool_output_unrecognized_non_blocking_response_falls_back_to_original_result() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, _request_rx) = mpsc::channel(); - let response_body = json!({ - "id": "chatcmpl-tool-output-missing-result", - "object": "chat.completion", - "created": 1, - "model": "", - "choices": [{ - "index": 0, - "message": { - "role": "assistant", - "content": "{\"tool_name\":\"weather_lookup\",\"arguments\":{\"city\":\"Phoenix\"}}" - }, - "finish_reason": "stop" - }], - "guardrails": { - "config_id": "safety-default", - "log": { - "activated_rails": [] - } - } - }) - .to_string(); - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "input": false, - "output": false, - "tool_output": true, - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - } - }))) - .await - .unwrap(); - - let result = tool_call_execute( - ToolCallExecuteParams::builder() - .name("weather_lookup") - .args(json!({"city": "Phoenix"})) - .func(Arc::new(move |_args| { - Box::pin(async move { Ok(json!({"forecast": "sunny"}).into()) }) - })) - .build(), - ) - .await - .unwrap(); - - assert_eq!(result.result, json!({"forecast": "sunny"})); -} - -#[tokio::test] -async fn remote_tool_output_does_not_run_when_tool_callback_errors() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "input": false, - "output": false, - "tool_output": true, - "remote": { - "endpoint": unused_local_endpoint(), - "config_id": "safety-default" - } - }))) - .await - .unwrap(); - - let error = tool_call_execute( - ToolCallExecuteParams::builder() - .name("weather_lookup") - .args(json!({"city": "Phoenix"})) - .func(Arc::new(move |_args| { - Box::pin(async move { - Err(crate::error::FlowError::Internal( - "tool callback failed".to_string(), - )) - }) - })) - .build(), - ) - .await - .unwrap_err(); - - match error { - crate::error::FlowError::Internal(message) => { - assert_eq!(message, "tool callback failed"); - } - other => panic!("unexpected error: {other}"), - } -} - -#[tokio::test] -async fn remote_tool_checks_forward_context_state_and_thread_id() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let (request_tx, request_rx) = mpsc::channel(); - let response_body = json!({ - "id": "chatcmpl-tool-output-context", - "object": "chat.completion", - "created": 1, - "model": "", - "choices": [{ - "index": 0, - "message": { - "role": "tool", - "name": "weather_lookup", - "tool_call_id": "call_weather_lookup_1", - "content": "{\"forecast\":\"cloudy\"}" - }, - "finish_reason": "stop" - }], - "guardrails": { - "config_id": "safety-default", - "log": { - "activated_rails": [] - } - } - }) - .to_string(); - let http_response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", - response_body.len(), - response_body - ) - .into_bytes(); - spawn_http_responder(listener, http_response, request_tx); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "input": false, - "output": false, - "tool_output": true, - "remote": { - "endpoint": format!("http://{address}"), - "config_id": "safety-default" - }, - "request_defaults": { - "context": {"tenant": "smoke"}, - "thread_id": "1234567890abcdef", - "state": {"events": []} - } - }))) - .await - .unwrap(); - - let result = tool_call_execute( - ToolCallExecuteParams::builder() - .name("weather_lookup") - .args(json!({"city": "Phoenix"})) - .func(Arc::new(move |_args| { - Box::pin(async move { Ok(json!({"forecast": "sunny"}).into()) }) - })) - .build(), - ) - .await - .unwrap(); - - assert_eq!(result.result, json!({"forecast": "cloudy"})); - - let captured = recv_captured_request(&request_rx); - let request_json: Json = serde_json::from_slice(&captured.body).unwrap(); - assert_eq!( - request_json["guardrails"]["context"], - json!({"tenant": "smoke"}) - ); - assert_eq!( - request_json["guardrails"]["thread_id"], - json!("1234567890abcdef") - ); - assert_eq!(request_json["guardrails"]["state"], json!({"events": []})); -} - -#[tokio::test] -async fn remote_tool_only_configuration_does_not_intercept_llm_calls() { - let _guard = crate::plugins::nemo_guardrails::test_mutex() - .lock() - .unwrap_or_else(|err| err.into_inner()); - reset_runtime(); - setup_isolated_thread(); - - test_initialize_plugin_host_exact(plugin_config(json!({ - "mode": "remote", - "input": false, - "output": false, - "tool_output": true, - "remote": { - "endpoint": unused_local_endpoint(), - "config_id": "safety-default" - } - }))) - .await - .unwrap(); - - let expected = json!({"response": "original"}); - let func: LlmExecutionNextFn = Arc::new(move |_req| { - let expected = expected.clone(); - Box::pin(async move { Ok(expected) }) - }); - - let response = llm_call_execute( - LlmCallExecuteParams::builder() - .name("openai") - .request(make_chat_request(false)) - .func(func) - .attributes(LlmAttributes::empty()) - .response_codec(Arc::new(OpenAIChatCodec) as Arc) - .build(), - ) - .await - .unwrap(); - - assert_eq!(response, json!({"response": "original"})); -} diff --git a/crates/node/Cargo.toml b/crates/node/Cargo.toml index 8d6faca81e..e64b4ffd8d 100644 --- a/crates/node/Cargo.toml +++ b/crates/node/Cargo.toml @@ -21,7 +21,7 @@ test = false __skip-implicit-config = ["nemo-relay/__skip-implicit-config"] [dependencies] -nemo-relay = { workspace = true, features = ["atof-streaming", "guardrails-remote", "object-store", "worker-grpc"] } +nemo-relay = { workspace = true, features = ["atof-streaming", "object-store", "worker-grpc"] } nemo-relay-adaptive = { workspace = true, features = ["redis-backend"] } nemo-relay-pii-redaction.workspace = true chrono = "0.4" diff --git a/crates/python/Cargo.toml b/crates/python/Cargo.toml index 9ca47ccddf..560ea83a85 100644 --- a/crates/python/Cargo.toml +++ b/crates/python/Cargo.toml @@ -21,7 +21,7 @@ crate-type = ["cdylib", "rlib"] __skip-implicit-config = ["nemo-relay/__skip-implicit-config"] [dependencies] -nemo-relay = { workspace = true, features = ["atof-streaming", "guardrails-remote", "object-store", "worker-grpc"] } +nemo-relay = { workspace = true, features = ["atof-streaming", "object-store", "worker-grpc"] } nemo-relay-adaptive = { workspace = true, features = ["redis-backend"] } nemo-relay-pii-redaction.workspace = true pyo3 = { version = "0.29.0", features = ["abi3", "abi3-py311", "experimental-inspect", "macros"] } diff --git a/crates/python/src/lib.rs b/crates/python/src/lib.rs index a5f7ed04bf..aed0ddfca9 100644 --- a/crates/python/src/lib.rs +++ b/crates/python/src/lib.rs @@ -75,7 +75,3 @@ fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> { #[cfg(test)] #[path = "../tests/coverage/coverage_tests.rs"] mod coverage_tests; - -#[cfg(test)] -#[path = "../tests/coverage/nemo_guardrails_coverage_tests.rs"] -mod nemo_guardrails_coverage_tests; diff --git a/crates/python/tests/coverage/nemo_guardrails_coverage_tests.rs b/crates/python/tests/coverage/nemo_guardrails_coverage_tests.rs deleted file mode 100644 index fd1d4b55ef..0000000000 --- a/crates/python/tests/coverage/nemo_guardrails_coverage_tests.rs +++ /dev/null @@ -1,921 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Coverage tests for Python-facing local NeMo Guardrails integration. - -use std::ffi::CString; -use std::fs; -use std::panic::{AssertUnwindSafe, catch_unwind}; -use std::path::PathBuf; -use std::process::{self, Command, Stdio}; -use std::sync::{ - Mutex, - atomic::{AtomicUsize, Ordering}, -}; - -use nemo_relay::api::runtime::{NemoRelayContextState, global_context}; -use nemo_relay::plugin::{ - PluginComponentSpec, PluginConfig, test_close_plugin_host, test_initialize_plugin_host_exact, -}; -use pyo3::prelude::*; -use pyo3::types::{PyDict, PyModule}; -use serde_json::json; - -static NEXT_FAKE_GUARDRAILS_ID: AtomicUsize = AtomicUsize::new(1); -static SERIAL_TEST_MUTEX: Mutex<()> = Mutex::new(()); - -fn load_module<'py>(py: Python<'py>, code: &str) -> Bound<'py, PyModule> { - let code = CString::new(code).unwrap(); - let file_name = CString::new("nemo_guardrails_coverage_tests.py").unwrap(); - let module_name = CString::new("nemo_guardrails_coverage_tests").unwrap(); - PyModule::from_code(py, &code, &file_name, &module_name).unwrap() -} - -fn python_package_dir() -> PathBuf { - PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../python") -} - -struct FakeGuardrailsPackage { - root: PathBuf, - module_name: String, - python_executable: PathBuf, -} - -impl FakeGuardrailsPackage { - fn new(py: Python<'_>, module_name: &str, version: &str, implementation: &str) -> Self { - let id = NEXT_FAKE_GUARDRAILS_ID.fetch_add(1, Ordering::Relaxed); - let root = std::env::temp_dir().join(format!( - "nemo_relay_python_fake_guardrails_{}_{}", - process::id(), - id - )); - let package = root.join(module_name); - fs::create_dir_all(package.join("rails/llm")).unwrap(); - fs::write(package.join("rails/__init__.py"), "").unwrap(); - fs::write(package.join("rails/llm/__init__.py"), "").unwrap(); - fs::write(package.join("rails/llm/options.py"), fake_options_module()).unwrap(); - fs::write( - package.join("__init__.py"), - fake_root_module(version, implementation), - ) - .unwrap(); - - let python_executable = PathBuf::from(python_executable_for_worker(py)); - - Self { - root, - module_name: module_name.to_string(), - python_executable, - } - } -} - -fn python_executable_for_worker(py: Python<'_>) -> String { - let sys_executable = py - .import("sys") - .and_then(|sys| sys.getattr("executable")) - .and_then(|executable| executable.extract::()) - .ok(); - - for executable in [ - std::env::var("PYO3_PYTHON").ok(), - std::env::var("UV_PYTHON").ok(), - sys_executable, - Some("python3".to_string()), - ] - .into_iter() - .flatten() - { - let executable = executable.trim(); - if executable.is_empty() { - continue; - } - if Command::new(executable) - .arg("-c") - .arg("import sys") - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .status() - .map(|status| status.success()) - .unwrap_or(false) - { - return executable.to_string(); - } - } - - "python3".to_string() -} - -impl Drop for FakeGuardrailsPackage { - fn drop(&mut self) { - let _ = fs::remove_dir_all(&self.root); - } -} - -fn fake_guardrails_module_prelude( - module_name: &str, - python_dir: &str, - python_executable: &str, - python_path: &str, -) -> String { - format!( - r#" -import sys - -sys.path.insert(0, {python_dir:?}) - -MODULE_NAME = {module_name:?} -PYTHON_EXECUTABLE = {python_executable:?} -PYTHONPATH = {python_path:?} -"#, - python_dir = python_dir, - module_name = module_name, - python_executable = python_executable, - python_path = python_path, - ) -} - -fn fake_options_module() -> &'static str { - r#" -class RailType: - INPUT = "input" - OUTPUT = "output" - -class RailStatus: - BLOCKED = "blocked" - MODIFIED = "modified" - PASSED = "passed" -"# -} - -fn fake_root_module(version: &str, implementation: &str) -> String { - format!( - r#" -import types -from .rails.llm.options import RailStatus - -__version__ = {version:?} - -class Result: - def __init__(self, status, content=None, rail=None): - self.status = status - self.content = content - self.rail = rail - -class RailsConfig: - @staticmethod - def from_content(*, colang_content=None, yaml_content=None): - return {{"yaml": yaml_content, "colang": colang_content}} - - @staticmethod - def from_path(path): - return {{"path": path}} - -{implementation} -"# - ) -} - -fn check_sequence_guardrails() -> &'static str { - r#" -class LLMRails: - def __init__(self, config): - self.config = config - self._check_results = [ - Result(RailStatus.MODIFIED, content="sanitized user"), - Result(RailStatus.BLOCKED, rail="output-policy"), - Result(RailStatus.MODIFIED, content='{"arguments": {"city": "Boston"}}'), - Result(RailStatus.MODIFIED, content='{"result": {"ok": true}}'), - ] - - async def check_async(self, messages, rail_types): - return self._check_results.pop(0) -"# -} - -fn tool_sequence_guardrails() -> &'static str { - r#" -class LLMRails: - def __init__(self, config): - self.config = config - self._check_results = [ - Result(RailStatus.MODIFIED, content="sanitized user"), - Result(RailStatus.PASSED), - Result(RailStatus.MODIFIED, content='{"arguments": {"city": "Boston"}}'), - Result(RailStatus.MODIFIED, content='{"result": {"ok": true}}'), - ] - - async def check_async(self, messages, rail_types): - return self._check_results.pop(0) -"# -} - -fn streaming_guardrails() -> &'static str { - r#" -class LLMRails: - def __init__(self, config): - yaml = str(config.get("yaml", "")) - stream_first = "stream_first_false" not in yaml - self.config = types.SimpleNamespace( - rails=types.SimpleNamespace( - output=types.SimpleNamespace( - flows=["self check output"], - streaming=types.SimpleNamespace(enabled=True, stream_first=stream_first), - ) - ) - ) - self._stream_calls = 0 - - async def check_async(self, messages, rail_types): - return Result(RailStatus.PASSED) - - def stream_async(self, *, messages=None, generator=None, include_metadata=False): - async def _run(): - self._stream_calls += 1 - call_index = self._stream_calls - async for chunk in generator: - if call_index == 1: - yield chunk - if call_index > 1: - yield '{"error": {"message": "Blocked by output rails: output-policy", "type": "guardrails_violation"}}' - return _run() -"# -} - -fn with_isolated_nemo_relay_modules( - py: Python<'_>, - native_module: &Bound<'_, PyModule>, - f: impl FnOnce() -> T, -) -> T { - let _serial_guard = SERIAL_TEST_MUTEX.lock().unwrap(); - let sys = py.import("sys").unwrap(); - let modules = sys - .getattr("modules") - .unwrap() - .cast_into::() - .unwrap(); - let saved_modules = modules - .iter() - .filter_map(|(name, module)| { - let name = name.extract::().ok()?; - if name == "nemo_relay" || name.starts_with("nemo_relay.") { - Some((name, module.unbind())) - } else { - None - } - }) - .collect::>(); - - clear_nemo_relay_modules(&modules); - modules - .set_item("nemo_relay._native", native_module.clone()) - .unwrap(); - - let result = catch_unwind(AssertUnwindSafe(f)); - - clear_nemo_relay_modules(&modules); - for (name, module) in saved_modules { - modules.set_item(name, module).unwrap(); - } - reset_runtime_state(); - - match result { - Ok(value) => value, - Err(payload) => std::panic::resume_unwind(payload), - } -} - -fn clear_nemo_relay_modules(modules: &Bound<'_, PyDict>) { - let module_names = modules - .iter() - .filter_map(|(name, _)| name.extract::().ok()) - .filter(|name| name == "nemo_relay" || name.starts_with("nemo_relay.")) - .collect::>(); - - for name in module_names { - modules.del_item(name).unwrap(); - } -} - -fn with_event_loop(py: Python<'_>, f: impl FnOnce(Bound<'_, PyAny>) -> T) -> T { - let asyncio = py.import("asyncio").unwrap(); - #[cfg(windows)] - { - let policy = asyncio - .getattr("WindowsSelectorEventLoopPolicy") - .unwrap() - .call0() - .unwrap(); - asyncio - .call_method1("set_event_loop_policy", (policy,)) - .unwrap(); - } - let event_loop = asyncio.call_method0("new_event_loop").unwrap(); - asyncio - .call_method1("set_event_loop", (&event_loop,)) - .unwrap(); - let result = catch_unwind(AssertUnwindSafe(|| f(event_loop.clone().into_any()))); - let drain = PyModule::from_code( - py, - &CString::new( - "import asyncio\nasync def drain(loop):\n current = asyncio.current_task(loop)\n pending = asyncio.all_tasks(loop) - {current}\n for task in pending:\n task.cancel()\n if pending:\n await asyncio.gather(*pending, return_exceptions=True)\n", - ) - .unwrap(), - &CString::new("drain_test_loop.py").unwrap(), - &CString::new("drain_test_loop").unwrap(), - ) - .unwrap() - .getattr("drain") - .unwrap() - .call1((&event_loop,)) - .unwrap(); - event_loop - .call_method1("run_until_complete", (drain,)) - .unwrap(); - asyncio - .call_method1("set_event_loop", (py.None(),)) - .unwrap(); - event_loop.call_method0("close").unwrap(); - #[cfg(windows)] - asyncio - .call_method1("set_event_loop_policy", (py.None(),)) - .unwrap(); - match result { - Ok(value) => value, - Err(payload) => std::panic::resume_unwind(payload), - } -} - -fn reset_runtime_state() { - let _ = test_close_plugin_host(); - let context = global_context(); - *context.write().unwrap() = NemoRelayContextState::new(); -} - -#[test] -fn test_native_pymodule_entrypoint_registers_bindings_without_local_provider_install() { - let _python = crate::test_support::init_python_test(); - let _serial_guard = SERIAL_TEST_MUTEX.lock().unwrap(); - reset_runtime_state(); - Python::attach(|py| { - let module = PyModule::new(py, "_native_guardrails_provider").unwrap(); - crate::_native(&module).unwrap(); - }); - - let runtime = tokio::runtime::Runtime::new().unwrap(); - let error = runtime - .block_on(test_initialize_plugin_host_exact(PluginConfig { - version: 1, - components: vec![PluginComponentSpec { - kind: "nemo_guardrails".to_string(), - enabled: true, - config: serde_json::from_value(json!({ - "mode": "local", - "codec": "openai_chat", - "config_path": "./rails" - })) - .unwrap(), - }], - policy: Default::default(), - })) - .unwrap_err(); - - reset_runtime_state(); - match error { - nemo_relay::plugin::PluginError::RegistrationFailed(message) => { - assert!( - message.contains( - "NeMo Guardrails is required for the built-in NeMo Guardrails local backend" - ), - "unexpected message: {message}" - ); - } - other => panic!("unexpected error: {other}"), - } -} - -#[test] -fn test_guardrails_local_runtime_enforces_llm_input_and_output_checks() { - let _python = crate::test_support::init_python_test(); - reset_runtime_state(); - - Python::attach(|py| { - let native_module = PyModule::new(py, "_native_guardrails_local_runtime").unwrap(); - crate::_native(&native_module).unwrap(); - - with_isolated_nemo_relay_modules(py, &native_module, || { - let fake = FakeGuardrailsPackage::new( - py, - "fake_guardrails_local_runtime", - "0.22.0", - check_sequence_guardrails(), - ); - let python_dir = python_package_dir(); - let prelude = fake_guardrails_module_prelude( - &fake.module_name, - &python_dir.display().to_string(), - &fake.python_executable.display().to_string(), - &fake.root.display().to_string(), - ); - let module = load_module( - py, - &format!( - r#" -{prelude} - -import nemo_relay - -async def run_case(): - stack = nemo_relay.create_scope_stack() - nemo_relay.set_thread_scope_stack(stack) - activation = await nemo_relay.plugin.initialize( - {{ - "version": 1, - "components": [ - {{ - "kind": "nemo_guardrails", - "enabled": True, - "config": {{ - "mode": "local", - "codec": "openai_chat", - "config_yaml": "models: []", - "input": True, - "output": True, - "tool_input": True, - "tool_output": True, - "local": {{ - "python_module": MODULE_NAME, - "python_executable": PYTHON_EXECUTABLE, - "python_path": PYTHONPATH, - }}, - }}, - }} - ], - }} - ) - - request = nemo_relay.LLMRequest( - {{}}, - {{ - "model": "gpt-4o-mini", - "messages": [{{"role": "user", "content": "unsafe"}}], - }}, - ) - seen_request_messages = [] - - async def next_call(req): - seen_request_messages.append(req.content["messages"][-1]["content"]) - return {{ - "choices": [{{"message": {{"role": "assistant", "content": "safe reply"}}}}], - "id": "resp_1", - "model": "gpt-4o-mini", - }} - - try: - await nemo_relay.llm.execute( - "demo", - request, - next_call, - response_codec=nemo_relay.codecs.OpenAIChatCodec(), - ) - except RuntimeError as error: - llm_error = str(error) - else: - raise AssertionError("expected output rail block") - - await activation.close() - return {{ - "llm_error": llm_error, - "seen_request_messages": seen_request_messages, - }} -"#, - prelude = prelude, - ), - ); - - let result_json = with_event_loop(py, |event_loop| { - let coroutine = module.getattr("run_case").unwrap().call0().unwrap(); - let result = event_loop - .call_method1("run_until_complete", (coroutine,)) - .unwrap(); - crate::convert::py_to_json(&result).unwrap() - }); - - assert_eq!( - result_json["seen_request_messages"][0], - json!("sanitized user") - ); - assert!( - result_json["llm_error"] - .as_str() - .unwrap() - .contains("output rail blocked the LLM call"), - "unexpected error: {}", - result_json["llm_error"] - ); - assert!( - result_json["llm_error"] - .as_str() - .unwrap() - .contains("output-policy"), - "unexpected error: {}", - result_json["llm_error"] - ); - }); - }); - - reset_runtime_state(); -} - -#[test] -fn test_guardrails_local_runtime_rejects_unsupported_nemoguardrails_version() { - let _python = crate::test_support::init_python_test(); - reset_runtime_state(); - - Python::attach(|py| { - let native_module = PyModule::new(py, "_native_guardrails_version").unwrap(); - crate::_native(&native_module).unwrap(); - - with_isolated_nemo_relay_modules(py, &native_module, || { - let fake = FakeGuardrailsPackage::new( - py, - "fake_guardrails_bad_version", - "0.21.0", - check_sequence_guardrails(), - ); - let python_dir = python_package_dir(); - let prelude = fake_guardrails_module_prelude( - &fake.module_name, - &python_dir.display().to_string(), - &fake.python_executable.display().to_string(), - &fake.root.display().to_string(), - ); - let module = load_module( - py, - &format!( - r#" -{prelude} - -import nemo_relay - -async def run_case(): - await nemo_relay.plugin.initialize( - {{ - "version": 1, - "components": [ - {{ - "kind": "nemo_guardrails", - "enabled": True, - "config": {{ - "mode": "local", - "codec": "openai_chat", - "config_yaml": "models: []", - "input": True, - "local": {{ - "python_module": MODULE_NAME, - "python_executable": PYTHON_EXECUTABLE, - "python_path": PYTHONPATH, - }}, - }}, - }} - ], - }} - ) -"#, - prelude = prelude, - ), - ); - - let error = with_event_loop(py, |event_loop| { - let coroutine = module.getattr("run_case").unwrap().call0().unwrap(); - event_loop - .call_method1("run_until_complete", (coroutine,)) - .unwrap_err() - .to_string() - }); - - assert!( - error.contains("requires nemoguardrails==0.22.0"), - "unexpected error: {error}" - ); - assert!(error.contains("0.21.0"), "unexpected error: {error}"); - }); - }); - - reset_runtime_state(); -} - -#[test] -fn test_guardrails_local_runtime_enforces_streamed_output_rails() { - let _python = crate::test_support::init_python_test(); - reset_runtime_state(); - - Python::attach(|py| { - let native_module = PyModule::new(py, "_native_guardrails_streaming").unwrap(); - crate::_native(&native_module).unwrap(); - - with_isolated_nemo_relay_modules(py, &native_module, || { - let fake = FakeGuardrailsPackage::new( - py, - "fake_guardrails_streaming", - "0.22.0", - streaming_guardrails(), - ); - let python_dir = python_package_dir(); - let prelude = fake_guardrails_module_prelude( - &fake.module_name, - &python_dir.display().to_string(), - &fake.python_executable.display().to_string(), - &fake.root.display().to_string(), - ); - let module = load_module( - py, - &format!( - r#" -{prelude} - -event_log = [] - -import nemo_relay - -def plugin_config(config_yaml="models: []"): - return {{ - "version": 1, - "components": [ - {{ - "kind": "nemo_guardrails", - "enabled": True, - "config": {{ - "mode": "local", - "codec": "openai_chat", - "config_yaml": config_yaml, - "input": False, - "output": True, - "local": {{ - "python_module": MODULE_NAME, - "python_executable": PYTHON_EXECUTABLE, - "python_path": PYTHONPATH, - }}, - }}, - }} - ], - }} - -async def run_stream(request): - collected = [] - - def next_call(req): - async def _stream(): - event_log.append("source:hello") - yield {{"choices": [{{"delta": {{"content": "hello"}}}}]}} - await __import__("asyncio").sleep(0.01) - event_log.append("source:world") - yield {{"choices": [{{"delta": {{"content": "world"}}}}]}} - return _stream() - - stream = await nemo_relay.llm.stream_execute( - "demo", - request, - next_call, - collected.append, - lambda: {{"chunks": collected}}, - response_codec=nemo_relay.codecs.OpenAIChatCodec(), - ) - chunks = [] - async for chunk in stream: - event_log.append(f"yield:{{chunk['choices'][0]['delta']['content']}}") - chunks.append(chunk) - return chunks - -async def run_case(): - stack = nemo_relay.create_scope_stack() - nemo_relay.set_thread_scope_stack(stack) - event_log.clear() - activation = await nemo_relay.plugin.initialize(plugin_config()) - - request = nemo_relay.LLMRequest( - {{}}, - {{ - "model": "gpt-4o-mini", - "messages": [{{"role": "user", "content": "hello"}}], - }}, - ) - - allowed_chunks = await run_stream(request) - - try: - await run_stream(request) - except RuntimeError as error: - blocked = str(error) - else: - raise AssertionError("expected streamed output block") - - await activation.close() - activation = await nemo_relay.plugin.initialize( - plugin_config("stream_first_false") - ) - try: - await run_stream(request) - except RuntimeError as error: - modified = str(error) - else: - raise AssertionError("expected stream_first=false error") - - await activation.close() - return {{ - "allowed_chunks": allowed_chunks, - "blocked": blocked, - "event_log": event_log, - "modified": modified, - }} -"#, - prelude = prelude, - ), - ); - - let result = with_event_loop(py, |event_loop| { - let coroutine = module.getattr("run_case").unwrap().call0().unwrap(); - let result = event_loop - .call_method1("run_until_complete", (coroutine,)) - .unwrap(); - crate::convert::py_to_json(&result).unwrap() - }); - assert_eq!( - result["allowed_chunks"], - json!([ - {"choices": [{"delta": {"content": "hello"}}]}, - {"choices": [{"delta": {"content": "world"}}]} - ]) - ); - let event_log = result["event_log"].as_array().unwrap(); - for expected in ["source:hello", "source:world", "yield:hello", "yield:world"] { - assert!( - event_log.iter().any(|event| event == expected), - "missing event {expected}: {event_log:?}" - ); - } - let source_hello = event_log - .iter() - .position(|event| event == "source:hello") - .unwrap(); - let source_world = event_log - .iter() - .position(|event| event == "source:world") - .unwrap(); - let yield_hello = event_log - .iter() - .position(|event| event == "yield:hello") - .unwrap(); - let yield_world = event_log - .iter() - .position(|event| event == "yield:world") - .unwrap(); - assert!(source_hello < yield_hello); - assert!(source_world < yield_world); - assert!(yield_hello < yield_world); - assert!( - result["blocked"] - .as_str() - .unwrap() - .contains("output rail blocked the LLM call") - ); - assert!( - result["modified"] - .as_str() - .unwrap() - .contains("stream_first = true") - ); - }); - }); - - reset_runtime_state(); -} - -#[test] -fn test_local_guardrails_provider_initializes_and_enforces_managed_core_calls() { - let _python = crate::test_support::init_python_test(); - reset_runtime_state(); - - Python::attach(|py| { - let native_module = PyModule::new(py, "_native_guardrails_e2e").unwrap(); - crate::_native(&native_module).unwrap(); - - with_isolated_nemo_relay_modules(py, &native_module, || { - let fake = FakeGuardrailsPackage::new( - py, - "fake_guardrails_local_e2e", - "0.22.0", - tool_sequence_guardrails(), - ); - let python_dir = python_package_dir(); - let prelude = fake_guardrails_module_prelude( - &fake.module_name, - &python_dir.display().to_string(), - &fake.python_executable.display().to_string(), - &fake.root.display().to_string(), - ); - let module = load_module( - py, - &format!( - r#" -{prelude} - -import nemo_relay - -async def run_case(): - stack = nemo_relay.create_scope_stack() - nemo_relay.set_thread_scope_stack(stack) - - activation = await nemo_relay.plugin.initialize( - {{ - "version": 1, - "components": [ - {{ - "kind": "nemo_guardrails", - "enabled": True, - "config": {{ - "mode": "local", - "codec": "openai_chat", - "config_yaml": "models: []", - "input": True, - "output": True, - "tool_input": True, - "tool_output": True, - "local": {{ - "python_module": MODULE_NAME, - "python_executable": PYTHON_EXECUTABLE, - "python_path": PYTHONPATH, - }}, - }}, - }} - ], - }} - ) - - request = nemo_relay.LLMRequest( - {{}}, - {{ - "model": "gpt-4o-mini", - "messages": [{{"role": "user", "content": "unsafe"}}], - }}, - ) - - seen_request_messages = [] - async def llm_impl(req): - seen_request_messages.append(req.content["messages"][-1]["content"]) - return {{ - "choices": [{{"message": {{"role": "assistant", "content": "safe reply"}}}}], - "id": "resp_1", - "model": req.content["model"], - }} - - llm_result = await nemo_relay.llm.execute( - "demo", - request, - llm_impl, - response_codec=nemo_relay.codecs.OpenAIChatCodec(), - ) - - seen_tool_args = [] - async def tool_impl(args): - seen_tool_args.append(args) - return nemo_relay.ToolExecutionResult({{"raw": True}}) - - tool_result = await nemo_relay.tools.execute("weather_lookup", {{"city": "Phoenix"}}, tool_impl) - await activation.close() - return {{ - "llm_result": llm_result, - "tool_result": tool_result.result, - "seen_request_messages": seen_request_messages, - "seen_tool_args": seen_tool_args, - }} -"#, - prelude = prelude, - ), - ); - let result_json = with_event_loop(py, |event_loop| { - let coroutine = module.getattr("run_case").unwrap().call0().unwrap(); - let result = event_loop - .call_method1("run_until_complete", (coroutine,)) - .unwrap(); - crate::convert::py_to_json(&result).unwrap() - }); - - assert_eq!( - result_json["llm_result"]["choices"][0]["message"]["content"], - json!("safe reply") - ); - assert_eq!(result_json["tool_result"], json!({ "ok": true })); - assert_eq!( - result_json["seen_request_messages"][0], - json!("sanitized user") - ); - assert_eq!( - result_json["seen_tool_args"][0], - json!({ "city": "Boston" }) - ); - }); - }); - - reset_runtime_state(); -} diff --git a/docs/about-nemo-relay/concepts/plugins.mdx b/docs/about-nemo-relay/concepts/plugins.mdx index 87a5c1f525..46cc73dfe3 100644 --- a/docs/about-nemo-relay/concepts/plugins.mdx +++ b/docs/about-nemo-relay/concepts/plugins.mdx @@ -185,8 +185,8 @@ in how their component kind becomes available to the host: ## Built-In Plugin Components -The core runtime registers the `observability`, `nemo_guardrails`, and -`pricing` components before lookup, validation, and initialization. The CLI and +The core runtime registers the `observability` and `pricing` components before +lookup, validation, and initialization. The CLI and the Python and Node.js bindings also register `adaptive` and `pii_redaction`. Direct Rust applications must register Adaptive and PII Redaction from their component crates before validating or initializing either kind: @@ -224,20 +224,6 @@ from the plugin namespace instead of exposed in public config. Detailed observability plugin configuration belongs in [Observability Configuration](/configure-plugins/observability/configuration). -### NeMo Guardrails - -The core crate also ships a built-in `nemo_guardrails` plugin component. The -built-in integration is deprecated and scheduled for removal in NeMo Relay -0.9. There is no replacement in NeMo Relay 0.8; any replacement will target 0.9 -or later. Do not use the built-in component for new deployments. - -The current user-facing paths are the remote backend for Guardrails-service -integration and the Python-backed local backend that runs `nemoguardrails` -through a subprocess worker. - -Detailed Guardrails plugin configuration belongs in -[NeMo Guardrails Configuration](/configure-plugins/nemo-guardrails/configuration). - ### PII Redaction The `pii_redaction` component sanitizes emitted observability payloads without diff --git a/docs/about-nemo-relay/release-notes/index.mdx b/docs/about-nemo-relay/release-notes/index.mdx index eec26498c7..8acf6b1c69 100644 --- a/docs/about-nemo-relay/release-notes/index.mdx +++ b/docs/about-nemo-relay/release-notes/index.mdx @@ -171,6 +171,22 @@ instead of being forwarded. Malformed or repeated credential headers still return `401`. Refer to [Requests Without a Bound Credential](/daemon/configuration#requests-without-a-bound-credential). +### Compatibility and Migration + + + +The deprecated built-in `nemo_guardrails` component has been removed. Delete +legacy component entries before upgrading, including disabled entries. The +local and remote backends, public Rust configuration types, CLI editor support, +and `guardrails-remote` Cargo feature are no longer available. + + + +Future first-party NeMo Guardrails integration will be distributed separately +as a dynamic plugin and will not reuse the removed component configuration. +Refer to the +[Migration Guides](/reference/migration-guides#remove-the-built-in-nemo-guardrails-component). + ### LLM Execution Codec Context Relay now passes a codec context to every large language model (LLM) execution @@ -351,11 +367,6 @@ The following list covers unresolved issues from 0.9 and new limitations in - The Node.js binding and package workflows require Node.js 24 or later. - OpenClaw provides telemetry through public hooks. Its security and optimization coverage is partial because Relay does not manage its execution. -- The built-in `nemo_guardrails` plugin remains deprecated but is still - available in 0.10.0. No replacement is included. The remote backend - depends on its configured service's availability, latency, and policy behavior. - The local backend requires Python 3.11 or later and `nemoguardrails==0.22.0`. - Refer to [NeMo Guardrails](/configure-plugins/nemo-guardrails/about). - The personally identifiable information (PII) redaction plugin supports its local backend based on fixed rules. Local-model backend configuration is reserved for future work. diff --git a/docs/configure-plugins/about.mdx b/docs/configure-plugins/about.mdx index 2dccff5bd8..ccdb1d6380 100644 --- a/docs/configure-plugins/about.mdx +++ b/docs/configure-plugins/about.mdx @@ -30,9 +30,6 @@ entries in `plugins.toml`. or typed OpenTelemetry data. - [Adaptive](/configure-plugins/adaptive/about) configures adaptive runtime behavior. -- [NeMo Guardrails (Deprecated)](/configure-plugins/nemo-guardrails/about) - installs Guardrails-backed policy checks. The built-in plugin is scheduled - for removal in NeMo Relay 0.9 and should not be used for new deployments. - [PII Redaction](/configure-plugins/pii-redaction/about) sanitizes sensitive data in observability payloads. - [Model Pricing](/configure-plugins/model-pricing) configures catalog sources diff --git a/docs/configure-plugins/nemo-guardrails/about.mdx b/docs/configure-plugins/nemo-guardrails/about.mdx deleted file mode 100644 index 64c33088cb..0000000000 --- a/docs/configure-plugins/nemo-guardrails/about.mdx +++ /dev/null @@ -1,123 +0,0 @@ ---- -title: "NeMo Guardrails (Deprecated)" -sidebar-title: "About" -description: "Review the deprecated built-in NeMo Guardrails plugin for managed LLM and tool execution." -position: 1 ---- -{/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 */} - - -The built-in `nemo_guardrails` plugin is deprecated and scheduled for removal -in NeMo Relay 0.9. It remains available in NeMo Relay 0.8. There is no -replacement in NeMo Relay 0.8; any replacement will target 0.9 or later. - - - -Do not start a new deployment on the built-in plugin. Use this documentation to -maintain or remove an existing NeMo Relay 0.8 configuration. - -The built-in plugin component has kind `nemo_guardrails` and is available as a -first-party NeMo Relay plugin. - -The plugin supports these backend modes: - -- `remote` - - Calls a Guardrails service over HTTP(S), including streaming over the same - remote contract. -- `local` - - Calls `nemoguardrails` through a local `python3` worker subprocess instead - of a separate Guardrails service. - -## Existing Uses - -Existing NeMo Relay 0.8 deployments may use the built-in plugin to: - -- Apply Guardrails input and output checks around managed `llm.execute(...)` - calls. -- Apply Guardrails policy around managed tool execution. -- Configure Guardrails behavior through the same plugin config surface used by - other first-party NeMo Relay components. -- Keep Guardrails policy authoring in Guardrails-native config while NeMo Relay - owns when those checks run around managed execution. - -## Current Scope - -The built-in plugin currently exposes two user-facing modes: - -- `remote` for Guardrails-service integration over HTTP(S) -- `local` for `nemoguardrails` integration through a local Python worker - -Both modes support managed LLM `input` and `output`. The current mode-specific -differences are: - -- `remote` supports `request_defaults` pass-through but does not support managed - `tool_input` -- `local` supports managed `tool_input` and broader LLM codec coverage, but it - does not support `request_defaults` - -The `local` backend requires a `python3 >= 3.11` executable that can import -`nemoguardrails==0.22.0`. It does not embed Python into the NeMo Relay binary. - -## Managed Surfaces Versus Request Defaults - -Both `remote` mode and `local` mode share the same top-level plugin model, but -they do not implement every part of that model in the same way. - -At the plugin-model level, NeMo Guardrails uses two different concepts: - -- Top-level managed NeMo Relay execution surfaces: - - `input` - - `output` - - `tool_input` - - `tool_output` -- Guardrails backend request defaults: - - `request_defaults.context` - - `request_defaults.thread_id` - - `request_defaults.state` - - `request_defaults.rails` - - `request_defaults.llm_params` - - `request_defaults.llm_output` - - `request_defaults.output_vars` - - `request_defaults.log` - -This distinction matters: - -- Managed surfaces wrap real NeMo Relay execution boundaries such as - `llm.execute(...)` and `tools.execute(...)`. -- Managed surfaces give NeMo Relay an owned enforcement point around a known - runtime step. Depending on the backend and surface, Relay can block work, - allow it, or apply managed request or result handling before the application - sees the outcome. -- Managed surfaces also give NeMo Relay a stable runtime boundary for its own - middleware ordering, lifecycle behavior, and observability marks. - -The forwarded request-default side is more mode-specific: - -- In `remote` mode, `request_defaults` fields are forwarded to the selected - Guardrails backend as request semantics. They do not create new NeMo - Relay-native execution surfaces. -- In `local` mode, `request_defaults` is rejected instead of passed through. - -The overlap in names is important in `remote` mode: - -- Top-level `input` is a managed NeMo Relay execution surface. -- `request_defaults.rails.input` is a backend pass-through option. -- Top-level `output` is a managed NeMo Relay execution surface. -- `request_defaults.rails.output` is a backend pass-through option. -- Top-level `tool_input` is a managed NeMo Relay execution surface in the - plugin contract. The current stock-remote backend rejects it, while the local - backend supports it. -- `request_defaults.rails.tool_input` is a backend pass-through option. -- Top-level `tool_output` is a managed NeMo Relay execution surface. -- `request_defaults.rails.tool_output` is a backend pass-through option. - -In particular, `request_defaults.rails.dialog` and -`request_defaults.rails.retrieval` are pass-through options. They are not -separate managed middleware surfaces in NeMo Relay. - -## Pages - -- [NeMo Guardrails Configuration](/configure-plugins/nemo-guardrails/configuration) - documents the built-in component shape, mode boundaries, and the detailed - support matrix. diff --git a/docs/configure-plugins/nemo-guardrails/configuration.mdx b/docs/configure-plugins/nemo-guardrails/configuration.mdx deleted file mode 100644 index c0bd78cd89..0000000000 --- a/docs/configure-plugins/nemo-guardrails/configuration.mdx +++ /dev/null @@ -1,392 +0,0 @@ ---- -title: "NeMo Guardrails Configuration" -sidebar-title: "Configuration" -description: "Configure the deprecated built-in remote and local NeMo Guardrails backends." -position: 2 ---- -{/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 */} - - -The built-in `nemo_guardrails` plugin is deprecated and scheduled for removal -in NeMo Relay 0.9. Existing local and remote configurations continue to work in -NeMo Relay 0.8. There is no replacement in NeMo Relay 0.8; any replacement -will target 0.9 or later. - - - -Use this page only to maintain or remove an existing NeMo Relay 0.8 -configuration. Do not start a new deployment on the built-in plugin. The -component kind is `nemo_guardrails`. - -For plugin file discovery, precedence, merge behavior, editor controls, and -gateway conflict rules, refer to -[Plugin Configuration Files](/configure-plugins/plugin-configuration-files). - - -NeMo Relay plugin configuration uses the generic plugin document shape, so -field names stay `snake_case` in every binding and in `plugins.toml`. - - - -## Component Shape - -The top-level NeMo Guardrails object contains: - -| Field | Purpose | -|---|---| -| `version` | Guardrails config schema version. Defaults to `1`. | -| `mode` | Backend mode. Current values are `remote` and `local`. | -| `config_path` | Local-mode native Guardrails config directory path. | -| `config_yaml` | Local-mode inline native Guardrails YAML config. | -| `colang_content` | Optional inline Colang content for local mode when `config_yaml` is used. | -| `codec` | Managed LLM provider codec. | -| `input` | Enables managed LLM input checks. | -| `output` | Enables managed LLM output checks. | -| `tool_input` | Enables managed tool-argument checks before execution. | -| `tool_output` | Enables managed tool-result checks after execution. | -| `priority` | Middleware priority for installed execution intercepts. | -| `remote` | Remote backend settings. | -| `local` | Local backend settings. | -| `request_defaults` | Default request-time Guardrails semantics passed to the remote backend. | -| `policy` | Component-local handling for unknown fields and unsupported values. | - -At least one managed Guardrails surface must be enabled. - -## Backend Support - -The following table compares remote and local backend support: - -| Area | `remote` | `local` | -|---|---|---| -| Built-in component kind and config validation | Supported | Supported | -| Managed LLM `input` | Supported | Supported | -| Managed LLM `output` | Supported | Supported | -| Managed streaming LLM execution | Supported over the remote HTTP(S) contract | Supported; refer to [Streaming Boundary](#streaming-boundary) | -| Managed `tool_input` | Not supported against the stock Guardrails remote contract | Supported | -| Managed `tool_output` | Supported | Supported | -| `request_defaults` pass-through | Supported | Not supported | -| Codec support | `openai_chat` | `openai_chat`, `openai_responses`, `anthropic_messages`, `oci_genai`, `gemini_generate_content` | -| Runtime availability | Any runtime that includes the remote backend | Runtimes that can start `python3 >= 3.11` with `nemoguardrails==0.22.0` installed | - -## Remote Mode - -Use `remote` mode when NeMo Relay should call a Guardrails service over -HTTP(S), especially when Guardrails must be shared across runtimes, used from -non-Python environments, or deployed independently from the application -process. - -### Requirements - -To use `mode = "remote"`, the configured `remote.endpoint` must point at a -Guardrails service that NeMo Relay can reach from the running process and that -exposes the Guardrails remote HTTP(S) contract. - -The NeMo Relay plugin config activates Guardrails integration, but the -Guardrails service still owns the actual policy content. In practice, NeMo -Relay decides when managed checks run, while the Guardrails config decides what -to block, allow, or rewrite. - -### Remote Settings - -The remote backend accepts the following settings: - -| Field | Default | Notes | -|---|---|---| -| `remote.endpoint` | Required | Base `http://` or `https://` URL for the reachable Guardrails service. | -| `remote.config_id` | Omitted | One Guardrails configuration identifier. Set this or `remote.config_ids`, but not both. | -| `remote.config_ids` | `[]` | Multiple Guardrails configuration identifiers to combine. Set this or `remote.config_id`, but not both. | -| `remote.headers` | `{}` | Static string headers sent with every remote request. | -| `remote.timeout_millis` | `3000` | Positive request timeout in milliseconds. | - -### `plugins.toml` Example - -You can write this config directly in `plugins.toml`, or create and edit it -through the CLI with `nemo-relay plugins edit`. For plugin file discovery, -precedence, merge behavior, and editor controls, refer to -[Plugin Configuration Files](/configure-plugins/plugin-configuration-files). - -```toml -version = 1 - -[[components]] -kind = "nemo_guardrails" -enabled = true - -[components.config] -version = 1 -mode = "remote" -codec = "openai_chat" -input = true -output = true -tool_output = true - -[components.config.remote] -endpoint = "http://127.0.0.1:8000" -config_id = "live-smoke" -timeout_millis = 3000 - -[components.config.request_defaults.context] -tenant = "demo" - -[components.config.request_defaults.rails] -input = true -output = true - -[components.config.policy] -unknown_component = "warn" -unknown_field = "warn" -unsupported_value = "error" -``` - -This example configures the built-in remote mode for a Guardrails service that -uses `codec = "openai_chat"`, managed LLM `input` and `output`, managed -`tool_output`, and request-default pass-through for backend context plus -backend `input` and `output` rail selection. - -### Rules - -When `mode = "remote"`: - -- `remote.endpoint` is required. -- Exactly one of `remote.config_id` or `remote.config_ids` is required. -- `config_path`, `config_yaml`, and `colang_content` cannot be present. -- `local` settings cannot be present. -- The backend uses the Guardrails remote HTTP(S) contract for both non-streaming - and streaming LLM execution. - -### Codec Boundary - -The current built-in remote mode supports managed LLM execution only with: - -- `openai_chat` - -### Managed Tool Boundary - -The current remote mode supports managed `tool_output`. - -The current remote mode rejects managed `tool_input` explicitly because the -stock Guardrails remote contract does not activate pre-execution tool-call -rails from externally submitted `/v1/chat/completions` history. NeMo Relay -rejects `tool_input` in remote mode rather than leaving a silent -non-enforcing path. - -### Request Defaults - -`request_defaults` lets the built-in plugin pass request-time semantics through -to the selected remote backend. - -Supported request-default fields are: - -- `context` -- `thread_id` -- `state` -- `rails` -- `llm_params` -- `llm_output` -- `output_vars` -- `log` - -These are backend request options, not additional NeMo Relay-managed execution -surfaces. - -This includes fields whose names overlap with top-level managed surfaces: - -| Field | Meaning | -|---|---| -| Top-level `input` | Managed NeMo Relay LLM input surface | -| `request_defaults.rails.input` | Backend pass-through rail selection | -| Top-level `output` | Managed NeMo Relay LLM output surface | -| `request_defaults.rails.output` | Backend pass-through rail selection | -| Top-level `tool_input` | Managed NeMo Relay tool-input surface in the plugin model; not supported by the current stock-remote backend | -| `request_defaults.rails.tool_input` | Backend pass-through rail selection | -| Top-level `tool_output` | Managed NeMo Relay tool-output surface | -| `request_defaults.rails.tool_output` | Backend pass-through rail selection | - -The `rails` section can include: - -- `input` -- `output` -- `retrieval` -- `dialog` -- `tool_output` -- `tool_input` - -Those values are forwarded to the remote backend as request semantics. They do -not mean NeMo Relay owns separate managed retrieval or dialog execution -surfaces. `dialog` and `retrieval` are pass-through request options only. -Likewise, `request_defaults.rails.tool_input` is only a backend pass-through -selector. It does not make managed remote `tool_input` supported in the -stock-remote lane. - -For more targeted request-time pass-through, the remote backend also forwards -selectors like these: - -```toml -[components.config.request_defaults.rails] -input = true -output = true -retrieval = ["retrieve_relevant_chunks"] -dialog = true -tool_output = ["validate_tool_output"] -``` - -### Observability - -The current remote backend emits coarse backend-level marks for remote -Guardrails activity: - -- `nemo_guardrails.remote.start` -- `nemo_guardrails.remote.end` -- `nemo_guardrails.remote.error` - -## Local Mode - -Use `local` mode when NeMo Relay should call `nemoguardrails` through a local -Python worker subprocess instead of a separate Guardrails service. - -### Requirements - -To use `mode = "local"`, NeMo Relay must be able to start a `python3 >= 3.11` -executable that can import `nemoguardrails==0.22.0`. - -The built-in local backend starts a Python worker process and sends Guardrails -checks over a JSON-lines protocol. Use it when the runtime has direct access to -the Python Guardrails dependency and configuration files rather than a separate -Guardrails service. Install the tested local-mode Guardrails dependency with -`pip install nemoguardrails==0.22.0`. - -The same ownership boundary still applies: - -- NeMo Relay decides when managed checks run. -- Guardrails-native config still decides what to block, allow, or rewrite. - -### Local Settings - -The local backend accepts the following settings: - -| Field | Default | Notes | -|---|---|---| -| `config_path` | Omitted | Native Guardrails configuration directory. Set this or `config_yaml`, but not both. | -| `config_yaml` | Omitted | Inline native Guardrails YAML. Set this or `config_path`, but not both. | -| `colang_content` | Omitted | Inline Colang content. Use only with `config_yaml`. | -| `local.python_module` | `nemoguardrails` | Module that the local worker imports. Set a custom module only when the runtime exposes Guardrails through another import path. | -| `local.python_executable` | `NEMO_RELAY_PYTHON`, otherwise `python3` | Python executable that starts the local worker. | -| `local.python_path` | Omitted | Path prepended to the worker subprocess `PYTHONPATH`. | - -### `plugins.toml` Example - -You can write this config directly in `plugins.toml`, or create and edit it -through the CLI with `nemo-relay plugins edit`. For plugin file discovery, -precedence, merge behavior, and editor controls, refer to -[Plugin Configuration Files](/configure-plugins/plugin-configuration-files). - -```toml -version = 1 - -[[components]] -kind = "nemo_guardrails" -enabled = true - -[components.config] -version = 1 -mode = "local" -codec = "openai_chat" -input = true -output = true -tool_input = true -tool_output = true -config_path = "./rails" - -[components.config.local] -python_executable = "python3" - -[components.config.policy] -unknown_component = "warn" -unknown_field = "warn" -unsupported_value = "error" -``` - -This example configures the built-in local mode for a runtime that can start -`python3`, import `nemoguardrails`, and read a native Guardrails config -directory from `./rails`. - -For example, the Guardrails-side policy can look like this: - -```yaml -rails: - input: - flows: - - self check input - output: - flows: - - self check output -``` - -This Guardrails-side config defines the policy logic. The NeMo Relay plugin -config decides when those checks run. - -### Rules - -When `mode = "local"`: - -- Exactly one of `config_path` or `config_yaml` is required. -- `colang_content` can only be used with `config_yaml`. -- `remote` settings cannot be present. -- `request_defaults` is rejected. -- `local.python_module` is optional and only needed when the runtime should - import the Guardrails dependency from a custom Python module path instead of - the default `nemoguardrails` package. -- `local.python_executable` is optional and defaults to the - `NEMO_RELAY_PYTHON` environment variable when set, otherwise `python3`. -- `local.python_path` is optional and is prepended to `PYTHONPATH` only for - the local Guardrails worker subprocess. - -### Codec Boundary - -The current built-in local mode supports managed LLM execution with: - -- `openai_chat` -- `openai_responses` -- `anthropic_messages` -- `oci_genai` -- `gemini_generate_content` - -### Managed Tool Boundary - -The current local mode supports both: - -- managed `tool_input` -- managed `tool_output` - -### Streaming Boundary - -The current local mode supports streaming LLM input checks before the stream -callback runs. - -When output rails are configured, the current local mode uses Guardrails-native -streaming output rails and lets provider chunks flow while the local output rail -monitor evaluates the streamed text. That requires `rails.output.streaming.enabled = true` -in the Guardrails config. - -Guardrails calls the main streaming-output switch -`rails.output.streaming.stream_first`. - -When `stream_first = true`, the current local mode uses pass-through-first -streaming semantics: - -- Provider chunks can reach the caller immediately. -- Guardrails evaluates the streamed text in parallel. -- If Guardrails later blocks the stream, the call fails after some chunks have - already reached the caller. - -The current local mode does not support `rails.output.streaming.stream_first = false` -yet. That mode would require Guardrails-first chunk reconstruction: - -- Guardrails would need to evaluate streamed text before it releases chunks to - the caller. -- The local backend would then need to convert Guardrails-approved text into - valid provider-shaped stream chunks. - -That guarded-text-to-provider-chunk adapter does not exist yet in the current -local backend. diff --git a/docs/configure-plugins/plugin-configuration-files.mdx b/docs/configure-plugins/plugin-configuration-files.mdx index abe92f6272..beb471004d 100644 --- a/docs/configure-plugins/plugin-configuration-files.mdx +++ b/docs/configure-plugins/plugin-configuration-files.mdx @@ -245,11 +245,8 @@ Repository-local `.nemo-relay/plugins.toml` files are not discovered. ## Gateway Editing Files -Use the interactive editor for Observability, Adaptive, PII Redaction, and -existing NeMo Guardrails configuration. The built-in `nemo_guardrails` -component is deprecated and scheduled for removal in NeMo Relay 0.9; do not -use the editor to start a new Guardrails deployment. The editor also updates -the configuration of manifest-backed dynamic plugins that `plugins add` has +Use the interactive editor for Observability, Adaptive, PII Redaction, and the +configuration of manifest-backed dynamic plugins that `plugins add` has registered: ```bash @@ -444,9 +441,8 @@ your `initialize` config applies. ## Explicit Defaults and Overrides The editor writes explicit defaults for edited Observability and Adaptive -sections. It writes NeMo Guardrails, PII Redaction, and dynamic-plugin fields -only when readers configure them. Use the NeMo Guardrails editor only to -maintain or remove an existing configuration. In a layered config model, +sections. It writes PII Redaction and dynamic-plugin fields only when readers +configure them. In a layered config model, omitting a field means "inherit a lower precedence value"; it does not mean "delete that value." Use the dedicated `nemo-relay model-pricing` commands to manage model-pricing catalog sources. @@ -526,6 +522,5 @@ Use the component guides for field-level configuration: - [Adaptive Configuration](/configure-plugins/adaptive/configuration) - [Adaptive Cache Governor (ACG)](/configure-plugins/adaptive/acg) - [Adaptive Hints](/configure-plugins/adaptive/adaptive-hints) -- [NeMo Guardrails Configuration](/configure-plugins/nemo-guardrails/configuration) - [PII Redaction Configuration](/configure-plugins/pii-redaction/configuration) - [Model Pricing](/configure-plugins/model-pricing) diff --git a/docs/getting-started/configuration.mdx b/docs/getting-started/configuration.mdx index bda6482723..6fe9c7f0f2 100644 --- a/docs/getting-started/configuration.mdx +++ b/docs/getting-started/configuration.mdx @@ -44,11 +44,6 @@ A plugin configuration has these parts: Start with [Language Binding Plugins](/build-plugins/language-binding/about) when you need reusable middleware, subscribers, or adaptive behavior. -Use [NeMo Guardrails Configuration](/configure-plugins/nemo-guardrails/configuration) -only to maintain or remove an existing NeMo Relay 0.8 `nemo_guardrails` -component. The built-in plugin is deprecated and scheduled for removal in NeMo -Relay 0.9; do not use it for a new deployment. - The `nemo-relay` CLI gateway reads plugin files named `plugins.toml`. Refer to [Plugin Configuration Files](/configure-plugins/plugin-configuration-files) for file locations, precedence, merge behavior, editor controls, and validation diff --git a/docs/index.yml b/docs/index.yml index 230fdad414..4be7787f2d 100644 --- a/docs/index.yml +++ b/docs/index.yml @@ -41,9 +41,6 @@ navigation: - folder: ./configure-plugins/adaptive title: "Adaptive" title-source: frontmatter - - folder: ./configure-plugins/nemo-guardrails - title: "NeMo Guardrails (Deprecated)" - title-source: frontmatter - folder: ./configure-plugins/pii-redaction title: "PII Redaction" title-source: frontmatter diff --git a/docs/reference/migration-guides.mdx b/docs/reference/migration-guides.mdx index e70b55b054..dca749d0d5 100644 --- a/docs/reference/migration-guides.mdx +++ b/docs/reference/migration-guides.mdx @@ -13,6 +13,37 @@ releases. Apply them when upgrading directly from 0.9.0. ## Upgrade to NeMo Relay 0.10.0 +### Remove the Built-In NeMo Guardrails Component + + + +NeMo Relay 0.10 removes the deprecated built-in `nemo_guardrails` component, +including its local and remote backends, public Rust configuration types, CLI +editor support, and `guardrails-remote` Cargo feature. + + + +Remove every legacy `[[components]]` entry with +`kind = "nemo_guardrails"`, including disabled entries. Relay rejects these +entries with migration guidance instead of silently ignoring them. + +No legacy setting is read by a replacement plugin: + +- Local settings such as `config_path`, `config_yaml`, `colang_content`, + `request_defaults`, and rail selectors require new plugin configuration. +- Remote settings such as `endpoint` and `config_id` have no automatic worker + equivalent. The removed backend delegated full generation to a Guardrails + service. Do not treat a remote check endpoint as the same behavior. +- Keep any Guardrails configuration files you still need outside the removed + `[[components]]` entry while planning a migration. + +Future first-party NeMo Guardrails integration will use Relay's separately +packaged dynamic-plugin model. It is not a drop-in configuration change. When +the replacement is released, review its supported surfaces and migrate the old +component settings to its new configuration. Follow releases and installation +guidance in +[NVIDIA NeMo Relay Plugins](https://github.com/NVIDIA/NeMo-Relay-Plugins). + ### Update Packages and Plugin SDKs Update the Relay CLI, application bindings, and plugin SDK dependencies to diff --git a/docs/resources/glossary.mdx b/docs/resources/glossary.mdx index e59b78ecfb..ba102fa4eb 100644 --- a/docs/resources/glossary.mdx +++ b/docs/resources/glossary.mdx @@ -320,12 +320,6 @@ the rest of the documentation can use them consistently. and compatible Relay version, and it is not sandboxed. Refer to [Native Dynamic Plugins](/build-plugins/native/about). -**NeMo Guardrails** -: NeMo Guardrails is the deprecated built-in `nemo_guardrails` plugin component - for first-party guardrail policy around managed tool and LLM execution. It - supports remote Guardrails-service and Python-backed local backends and is - scheduled for removal in NeMo Relay 0.9. - **Next Function** : The next function is the continuation passed to an execution intercept. The intercept calls `next` to run the next intercept or the original callback. An diff --git a/fern/docs.yml b/fern/docs.yml index a5b33af963..a3c317ad66 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -100,9 +100,13 @@ redirects: - source: /nemo/relay/adaptive-plugin/adaptive-hints destination: /nemo/relay/configure-plugins/adaptive/adaptive-hints - source: /nemo/relay/nemo-guardrails-plugin/about - destination: /nemo/relay/configure-plugins/nemo-guardrails/about + destination: /nemo/relay/reference/migration-guides - source: /nemo/relay/nemo-guardrails-plugin/configuration - destination: /nemo/relay/configure-plugins/nemo-guardrails/configuration + destination: /nemo/relay/reference/migration-guides +- source: /nemo/relay/configure-plugins/nemo-guardrails/about + destination: /nemo/relay/reference/migration-guides +- source: /nemo/relay/configure-plugins/nemo-guardrails/configuration + destination: /nemo/relay/reference/migration-guides - source: /nemo/relay/pii-redaction-plugin/about destination: /nemo/relay/configure-plugins/pii-redaction/about - source: /nemo/relay/pii-redaction-plugin/configuration @@ -110,7 +114,7 @@ redirects: # Legacy release documentation - source: /nemo/relay/v0.3.0/build-plugins/nemoguardrails - destination: /nemo/relay/configure-plugins/nemo-guardrails/about + destination: /nemo/relay/reference/migration-guides # Build plugin guides - source: /nemo/relay/build-plugins/basic-guide