diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index e9b7b47ac..5ab33bfd3 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 30f378ba9..4d969a6b0 100644 --- a/crates/cli/src/plugins/editor_model.rs +++ b/crates/cli/src/plugins/editor_model.rs @@ -17,20 +17,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)] @@ -47,7 +33,6 @@ pub(super) struct ComponentEditorState { pub(super) enum EditableComponent { Observability(Box>), Adaptive(Box>), - NemoGuardrails(Box>), PiiRedaction(Box>), } @@ -56,7 +41,6 @@ impl EditableComponent { match self { Self::Observability(_) => "Observability", Self::Adaptive(_) => "Adaptive", - Self::NemoGuardrails(_) => "NeMo Guardrails (Deprecated)", Self::PiiRedaction(_) => "PII Redaction", } } @@ -65,7 +49,6 @@ impl EditableComponent { match self { Self::Observability(_) => ObservabilityConfig::editor_schema().fields, Self::Adaptive(_) => AdaptiveConfig::editor_schema().fields, - Self::NemoGuardrails(_) => NeMoGuardrailsConfig::editor_schema().fields, Self::PiiRedaction(_) => PiiRedactionConfig::editor_schema().fields, } } @@ -74,7 +57,6 @@ impl EditableComponent { match self { Self::Observability(state) => state.enabled, Self::Adaptive(state) => state.enabled, - Self::NemoGuardrails(state) => state.enabled, Self::PiiRedaction(state) => state.enabled, } } @@ -83,7 +65,6 @@ impl EditableComponent { match self { Self::Observability(state) => state.toggle_enabled(), Self::Adaptive(state) => state.toggle_enabled(), - Self::NemoGuardrails(state) => state.toggle_enabled(), Self::PiiRedaction(state) => state.toggle_enabled(), } } @@ -92,7 +73,6 @@ impl EditableComponent { match self { Self::Observability(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), } } @@ -101,7 +81,6 @@ impl EditableComponent { match self { Self::Observability(state) => state.reset_enabled(), Self::Adaptive(state) => state.reset_enabled(), - Self::NemoGuardrails(state) => state.reset_enabled(), Self::PiiRedaction(state) => state.reset_enabled(), } } @@ -110,7 +89,6 @@ impl EditableComponent { match self { Self::Observability(state) => observability_summary(state), Self::Adaptive(state) => adaptive_summary(state), - Self::NemoGuardrails(state) => nemo_guardrails_summary(state), Self::PiiRedaction(state) => pii_redaction_summary(state), } } @@ -119,9 +97,6 @@ impl EditableComponent { match self { Self::Observability(state) => section_configured(&state.config, field), 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) } @@ -138,10 +113,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(); @@ -163,10 +134,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(); @@ -179,7 +146,6 @@ impl EditableComponent { match self { Self::Observability(state) => store_observability_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), } } @@ -207,7 +173,6 @@ pub(super) fn editable_components( let components = vec![ EditableComponent::Observability(Box::new(component_observability_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) @@ -403,12 +368,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> { @@ -447,22 +406,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, @@ -793,23 +736,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> { @@ -854,21 +780,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, @@ -1084,35 +995,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 e92d943f9..769830814 100644 --- a/crates/cli/src/plugins/prompt.rs +++ b/crates/cli/src/plugins/prompt.rs @@ -171,10 +171,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 7d98b0b9e..4dd07f4f4 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 553f29604..28ffc11ff 100644 --- a/crates/cli/tests/coverage/shared/plugins_tests.rs +++ b/crates/cli/tests/coverage/shared/plugins_tests.rs @@ -17,25 +17,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, -}; - fn write_editor_dynamic_manifest( dir: &Path, plugin_id: &str, @@ -117,54 +98,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!( @@ -324,64 +257,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(); @@ -534,11 +409,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" @@ -1056,63 +926,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(); @@ -1356,93 +1169,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 { @@ -1788,55 +1514,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(); @@ -2532,21 +2209,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 @@ -2566,16 +2236,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] @@ -2680,34 +2340,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 { @@ -2734,213 +2366,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 04bbacc92..cf0659cf3 100644 --- a/crates/cli/tests/coverage/shared/server_tests.rs +++ b/crates/cli/tests/coverage/shared/server_tests.rs @@ -2580,6 +2580,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()); @@ -2610,6 +2632,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( @@ -2659,6 +2691,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 0983acc58..dcdbc8db4 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 c79188d97..2ca40a767 100644 --- a/crates/core/src/plugin.rs +++ b/crates/core/src/plugin.rs @@ -1151,10 +1151,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| { @@ -1162,7 +1158,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, ] .iter() @@ -1180,7 +1175,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() } diff --git a/crates/core/src/plugins/mod.rs b/crates/core/src/plugins/mod.rs index d090caafc..47c77a096 100644 --- a/crates/core/src/plugins/mod.rs +++ b/crates/core/src/plugins/mod.rs @@ -4,4 +4,3 @@ //! First-party plugin implementations for NeMo Relay Core. pub mod model_pricing; -pub mod nemo_guardrails; 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 3505d2d63..000000000 --- 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 e1618836f..000000000 --- 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 6bb54022c..000000000 --- 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 5909d9a28..000000000 --- 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 b50f6f490..000000000 --- a/crates/core/src/plugins/nemo_guardrails/python.rs +++ /dev/null @@ -1,1487 +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, 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, 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 bf1f713e3..000000000 --- a/crates/core/src/plugins/nemo_guardrails/remote.rs +++ /dev/null @@ -1,1202 +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, _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, _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/plugins/nemo_guardrails/component_tests.rs b/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs deleted file mode 100644 index 99e578c82..000000000 --- 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 10dde622d..000000000 --- a/crates/core/tests/unit/plugins/nemo_guardrails/local_python_tests.rs +++ /dev/null @@ -1,1383 +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)); -} 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 c53caf7ce..000000000 --- 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 c3025bdf6..000000000 --- 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 8d6faca81..e64b4ffd8 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 9ca47ccdd..560ea83a8 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 a5f7ed04b..aed0ddfca 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 fd1d4b55e..000000000 --- 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 87a5c1f52..46cc73dfe 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 289bfc231..fbb191771 100644 --- a/docs/about-nemo-relay/release-notes/index.mdx +++ b/docs/about-nemo-relay/release-notes/index.mdx @@ -33,6 +33,22 @@ For the full release history, including individual pull requests, refer to NeMo Relay 0.10 is under development. This page will track its user-visible changes, compatibility updates, and fixed known issues. +### 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). + ## Known Issues in 0.10 No known issues have been recorded for NeMo Relay 0.10 yet. diff --git a/docs/configure-plugins/about.mdx b/docs/configure-plugins/about.mdx index b1aff92ec..70f8f67ea 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 64c33088c..000000000 --- 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 c0bd78cd8..000000000 --- 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 abe92f627..beb471004 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 0431a0d4e..fd584c6fa 100644 --- a/docs/getting-started/configuration.mdx +++ b/docs/getting-started/configuration.mdx @@ -36,11 +36,6 @@ Plugins use a structured plugin configuration with: 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 72275febd..55fa921af 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 6df03d769..4795cb142 100644 --- a/docs/reference/migration-guides.mdx +++ b/docs/reference/migration-guides.mdx @@ -11,6 +11,37 @@ upgrade actions as they are identified during the 0.10 development cycle. ## Upgrade to NeMo Relay 0.10 +### 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). + ### Refresh Codex Provider Routing After upgrading the CLI, refresh personal Relay-managed Codex installations: diff --git a/docs/resources/glossary.mdx b/docs/resources/glossary.mdx index e59b78ecf..ba102fa4e 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 a5b33af96..a3c317ad6 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 diff --git a/skills/nemo-relay-get-started/SKILL.md b/skills/nemo-relay-get-started/SKILL.md index 5ffdafd2b..8ff7c990f 100644 --- a/skills/nemo-relay-get-started/SKILL.md +++ b/skills/nemo-relay-get-started/SKILL.md @@ -101,7 +101,6 @@ establish the reusable plugin path. If Observability already produced the proof, ask what outcome matters next and recommend exactly one plugin: - Adaptive -> adaptive runtime behavior and optimization -- NeMo Guardrails -> policy checks around managed execution - PII Redaction -> sanitization of sensitive observability payloads - Model Pricing -> cost estimates for managed LLM responses diff --git a/skills/nemo-relay-get-started/evals/evals.json b/skills/nemo-relay-get-started/evals/evals.json index 057501174..3908509d4 100644 --- a/skills/nemo-relay-get-started/evals/evals.json +++ b/skills/nemo-relay-get-started/evals/evals.json @@ -182,7 +182,7 @@ "expected_behavior": [ "The agent recognizes that Observability established the first reusable plugin path during the default quick start", "The agent explains the instrument-once, change-behavior-through-plugins value of Relay", - "The agent maps Adaptive to optimization, NeMo Guardrails to policy, PII Redaction to sensitive telemetry, and Model Pricing to cost estimates", + "The agent maps Adaptive to optimization, PII Redaction to sensitive telemetry, and Model Pricing to cost estimates", "The agent asks or infers the user's end goal before recommending exactly one additional plugin", "The agent previews the smallest relevant plugin configuration and obtains confirmation before enabling it", "The agent recommends additional instrumentation only when the demonstrated boundary does not cover the real workflow", diff --git a/skills/nemo-relay-get-started/references/built-in-integrations-try-now.md b/skills/nemo-relay-get-started/references/built-in-integrations-try-now.md index 9ed0ff77c..e38276bae 100644 --- a/skills/nemo-relay-get-started/references/built-in-integrations-try-now.md +++ b/skills/nemo-relay-get-started/references/built-in-integrations-try-now.md @@ -74,8 +74,8 @@ After the first proof: - Preserve the maintained attachment boundary and recommend one goal-aligned plugin as the primary next step. If the trial used only temporary inspection, - configure plugin-managed Observability first; otherwise choose Adaptive, - NeMo Guardrails, PII Redaction, or Model Pricing based on the user's outcome. + configure plugin-managed Observability first; otherwise choose Adaptive, PII + Redaction, or Model Pricing based on the user's outcome. - Use the matching maintained integration guide for broader coverage. - Use `nemo-relay-instrument-calls` only for application-owned calls that remain outside the integration. diff --git a/skills/nemo-relay-get-started/references/cli-try-now.md b/skills/nemo-relay-get-started/references/cli-try-now.md index 129333962..ef5d64f3f 100644 --- a/skills/nemo-relay-get-started/references/cli-try-now.md +++ b/skills/nemo-relay-get-started/references/cli-try-now.md @@ -232,11 +232,10 @@ the coding-agent session is already instrumented, and later behavior can change through plugin configuration without reinstrumenting that boundary. Ask which outcome matters next and recommend one built-in plugin: Adaptive for -optimization, NeMo Guardrails for policy, PII Redaction for sensitive payloads, -or Model Pricing for cost estimates. Use the plugin overview to show the -smallest next configuration. Do not enable multiple plugins or extend -instrumentation unless the user requests it or the current boundary is -insufficient. +optimization, PII Redaction for sensitive payloads, or Model Pricing for cost +estimates. Use the plugin overview to show the smallest next configuration. Do +not enable multiple plugins or extend instrumentation unless the user requests +it or the current boundary is insufficient. ## Troubleshoot The Smallest Failed Boundary diff --git a/skills/nemo-relay-get-started/references/manual-language-try-now.md b/skills/nemo-relay-get-started/references/manual-language-try-now.md index 1759134c4..6d802476b 100644 --- a/skills/nemo-relay-get-started/references/manual-language-try-now.md +++ b/skills/nemo-relay-get-started/references/manual-language-try-now.md @@ -62,9 +62,9 @@ payloads. After this proof, preserve the demonstrated boundary and recommend one goal-aligned plugin as the primary next step. If the trial used only a short-lived subscriber, configure plugin-managed Observability first; otherwise -choose Adaptive, NeMo Guardrails, PII Redaction, or Model Pricing based on the -user's outcome. Use `nemo-relay-instrument-calls` only when the demonstrated -boundary does not yet cover the real application workflow. +choose Adaptive, PII Redaction, or Model Pricing based on the user's outcome. +Use `nemo-relay-instrument-calls` only when the demonstrated boundary does not +yet cover the real application workflow. For all supported languages, see the [Quick Start](https://docs.nvidia.com/nemo/relay/dev/getting-started/quick-start).