diff --git a/crates/cli/src/server/mod.rs b/crates/cli/src/server/mod.rs index d1b146e9c..41e2adc00 100644 --- a/crates/cli/src/server/mod.rs +++ b/crates/cli/src/server/mod.rs @@ -20,8 +20,9 @@ use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::{Json, Router}; use nemo_relay::plugin::dynamic::{ - DynamicPluginKind, NativePluginActivation, NativePluginLoadSpec, WorkerPluginActivation, - WorkerPluginLoadSpec, load_native_plugins, load_worker_plugins, + DynamicPluginCapability, DynamicPluginKind, DynamicPluginManifest, NativePluginActivation, + NativePluginLoadSpec, WorkerPluginActivation, WorkerPluginLoadSpec, load_native_plugins, + load_worker_plugins, }; use nemo_relay::plugin::{ PluginComponentSpec, PluginConfig, clear_plugin_configuration, initialize_plugins_exact, @@ -977,13 +978,6 @@ impl PluginActivation { { return Err(CliError::Config(error.to_string())); } - plugin_config - .components - .extend(dynamic_plugins.iter().map(|plugin| PluginComponentSpec { - kind: plugin.plugin_id.clone(), - enabled: true, - config: plugin.config.clone(), - })); for plugin in &dynamic_plugins { if let Some(snapshot) = plugin.activation_snapshot.as_ref() { snapshot.verify_current()?; @@ -1038,6 +1032,40 @@ impl PluginActivation { }) }) .collect::, CliError>>()?; + let mut policy_components = Vec::new(); + let mut regular_components = Vec::new(); + for plugin in &dynamic_plugins { + let manifest_ref = plugin + .activation_snapshot + .as_ref() + .map(|snapshot| snapshot.activation_manifest_ref()) + .or_else(|| plugin.manifest_ref.clone()) + .ok_or_else(|| { + CliError::Config(format!( + "dynamic plugin '{}' has no manifest_ref in lifecycle state", + plugin.plugin_id + )) + })?; + let (manifest, _) = DynamicPluginManifest::load_from_path(&manifest_ref) + .map_err(|error| CliError::Config(error.to_string()))?; + let component = PluginComponentSpec { + kind: plugin.plugin_id.clone(), + enabled: true, + config: plugin.config.clone(), + }; + if manifest + .capabilities + .items + .contains(&DynamicPluginCapability::ExportActivationPolicy) + { + policy_components.push(component); + } else { + regular_components.push(component); + } + } + policy_components.append(&mut plugin_config.components); + policy_components.append(&mut regular_components); + plugin_config.components = policy_components; let snapshots = dynamic_plugins .iter() .filter_map(|plugin| plugin.activation_snapshot.clone()) diff --git a/crates/core/src/api/export_activation.rs b/crates/core/src/api/export_activation.rs new file mode 100644 index 000000000..e44834dbd --- /dev/null +++ b/crates/core/src/api/export_activation.rs @@ -0,0 +1,339 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Activation-time hooks for plugin-managed export targets. + +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::{Arc, Mutex, RwLock}; +use std::time::Duration; + +use crate::error::{FlowError, Result}; +pub use nemo_relay_types::plugin::{ + ExportActivationDecision, ExportActivationPolicyConfig, ExportActivationRequest, + ExportActivationTargetKind, ExportTargetRegistration, MAX_EXPORT_ACTIVATION_TIMEOUT_MILLIS, + MIN_EXPORT_ACTIVATION_TIMEOUT_MILLIS, +}; + +/// Asynchronous callback registered by one export-activation policy provider. +pub type ExportActivationPolicyFn = Arc< + dyn Fn( + ExportActivationRequest, + ) -> Pin> + Send>> + + Send + + Sync, +>; + +/// Deferred constructor for one local or remote export target. +pub type ExportTargetActivationFn = + Arc Pin> + Send>> + Send + Sync>; + +struct PendingExportTarget { + qualified_id: String, + registration: ExportTargetRegistration, + activate: ExportTargetActivationFn, +} + +/// Policy callbacks and pending export targets owned by one plugin activation. +#[derive(Default)] +#[doc(hidden)] +pub struct ExportActivationRegistry { + policies: RwLock>, + targets: Mutex>, + active_targets: Mutex>, +} + +#[doc(hidden)] +pub type ExportActivationPolicyRegistry = ExportActivationRegistry; + +impl ExportActivationRegistry { + /// Registers one provider callback for this activation. + pub fn register(&self, provider: &str, callback: ExportActivationPolicyFn) -> Result<()> { + let mut policies = self.policies.write().map_err(|error| { + FlowError::Internal(format!( + "export activation policy registry lock poisoned: {error}" + )) + })?; + if policies.contains_key(provider) { + return Err(FlowError::AlreadyExists(provider.to_string())); + } + policies.insert(provider.to_string(), callback); + Ok(()) + } + + /// Removes one provider callback, returning whether it existed. + pub fn deregister(&self, provider: &str) -> Result { + self.policies + .write() + .map(|mut policies| policies.remove(provider).is_some()) + .map_err(|error| { + FlowError::Internal(format!( + "export activation policy registry lock poisoned: {error}" + )) + }) + } + + pub(crate) async fn evaluate( + &self, + provider: &str, + request: ExportActivationRequest, + ) -> Result { + let callback = self + .policies + .read() + .map_err(|error| { + FlowError::Internal(format!( + "export activation policy registry lock poisoned: {error}" + )) + })? + .get(provider) + .cloned() + .ok_or_else(|| FlowError::NotFound(provider.to_string()))?; + callback(request).await + } + + /// Adds one deferred exporter constructor under an activation-unique ID. + pub fn register_target( + &self, + qualified_id: String, + registration: ExportTargetRegistration, + activate: ExportTargetActivationFn, + ) -> Result<()> { + validate_target_registration(®istration)?; + let mut targets = self.targets.lock().map_err(|error| { + FlowError::Internal(format!("export target registry lock poisoned: {error}")) + })?; + if targets + .iter() + .any(|target| target.qualified_id == qualified_id) + { + return Err(FlowError::AlreadyExists(qualified_id)); + } + targets.push(PendingExportTarget { + qualified_id, + registration, + activate, + }); + Ok(()) + } + + /// Removes one pending exporter constructor, returning whether it existed. + pub fn deregister_target(&self, qualified_id: &str) -> Result { + let remove = |targets: &Mutex>| -> Result { + let mut targets = targets.lock().map_err(|error| { + FlowError::Internal(format!("export target registry lock poisoned: {error}")) + })?; + let original_len = targets.len(); + targets.retain(|target| target.qualified_id != qualified_id); + Ok(targets.len() != original_len) + }; + let pending = remove(&self.targets)?; + let active = remove(&self.active_targets)?; + Ok(pending || active) + } + + pub(crate) async fn activate_targets(&self) -> Result<()> { + let targets = { + let mut targets = self.targets.lock().map_err(|error| { + FlowError::Internal(format!("export target registry lock poisoned: {error}")) + })?; + std::mem::take(&mut *targets) + }; + for target in targets { + if self + .target_allowed(&target.qualified_id, &target.registration) + .await + { + (target.activate)().await?; + self.active_targets + .lock() + .map_err(|error| { + FlowError::Internal(format!( + "active export target registry lock poisoned: {error}" + )) + })? + .push(target); + } + } + Ok(()) + } + + async fn target_allowed( + &self, + qualified_id: &str, + registration: &ExportTargetRegistration, + ) -> bool { + let Some(policy) = ®istration.activation_policy else { + return true; + }; + let request = ExportActivationRequest { + target_kind: registration.target_kind.clone(), + config: policy.config.clone(), + }; + let outcome = tokio::time::timeout( + export_activation_timeout(policy), + self.evaluate(&policy.provider, request), + ) + .await; + let (allowed, reason) = match outcome { + Ok(Ok(ExportActivationDecision::Allow)) => (true, "allowed"), + Ok(Ok(ExportActivationDecision::Deny)) => (false, "denied"), + Ok(Err(FlowError::NotFound(_))) => (false, "provider_unavailable"), + Ok(Err(_)) => (false, "provider_error"), + Err(_) => (false, "timeout"), + }; + if !allowed { + log::warn!( + target: "nemo_relay.plugin", + event = "export_activation_policy_denied", + provider = policy.provider.as_str(), + target_kind = registration.target_kind.as_str(), + target_id = qualified_id, + reason; + "Export target suppressed by activation policy" + ); + } + allowed + } +} + +fn validate_target_registration(registration: &ExportTargetRegistration) -> Result<()> { + if registration.id.trim().is_empty() || registration.id.trim() != registration.id { + return Err(FlowError::InvalidArgument( + "export target id must be nonblank and have no surrounding whitespace".into(), + )); + } + if let Some(policy) = ®istration.activation_policy { + if policy.provider.trim().is_empty() || policy.provider.trim() != policy.provider { + return Err(FlowError::InvalidArgument( + "export activation provider must be nonblank and have no surrounding whitespace" + .into(), + )); + } + if !(MIN_EXPORT_ACTIVATION_TIMEOUT_MILLIS..=MAX_EXPORT_ACTIVATION_TIMEOUT_MILLIS) + .contains(&policy.timeout_millis) + { + return Err(FlowError::InvalidArgument(format!( + "export activation timeout_millis must be between {MIN_EXPORT_ACTIVATION_TIMEOUT_MILLIS} and {MAX_EXPORT_ACTIVATION_TIMEOUT_MILLIS}" + ))); + } + } + Ok(()) +} + +fn export_activation_timeout(policy: &ExportActivationPolicyConfig) -> Duration { + Duration::from_millis(policy.timeout_millis.clamp( + MIN_EXPORT_ACTIVATION_TIMEOUT_MILLIS, + MAX_EXPORT_ACTIVATION_TIMEOUT_MILLIS, + )) +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use serde_json::json; + + use super::*; + + fn registration(id: &str, provider: Option<&str>, allow: bool) -> ExportTargetRegistration { + ExportTargetRegistration { + id: id.into(), + target_kind: ExportActivationTargetKind::new("example.telemetry.exporter").unwrap(), + activation_policy: provider.map(|provider| ExportActivationPolicyConfig { + provider: provider.into(), + timeout_millis: 30_000, + config: json!({"allow": allow}), + }), + } + } + + #[tokio::test] + async fn self_provider_gates_deferred_targets_once_and_missing_provider_denies() { + let registry = ExportActivationRegistry::default(); + let policy_calls = Arc::new(AtomicUsize::new(0)); + let policy_calls_callback = Arc::clone(&policy_calls); + registry + .register( + "example.plugin", + Arc::new(move |request| { + let policy_calls = Arc::clone(&policy_calls_callback); + Box::pin(async move { + policy_calls.fetch_add(1, Ordering::SeqCst); + Ok(if request.config["allow"] == true { + ExportActivationDecision::Allow + } else { + ExportActivationDecision::Deny + }) + }) + }), + ) + .unwrap(); + registry + .register( + "error.plugin", + Arc::new(|_| Box::pin(async { Err(FlowError::Internal("policy failed".into())) })), + ) + .unwrap(); + + let activations = Arc::new(AtomicUsize::new(0)); + for (id, provider, allow) in [ + ("ungated", None, true), + ("allowed", Some("example.plugin"), true), + ("denied", Some("example.plugin"), false), + ("missing", Some("missing.plugin"), true), + ("error", Some("error.plugin"), true), + ] { + let activations = Arc::clone(&activations); + registry + .register_target( + format!("example:{id}"), + registration(id, provider, allow), + Arc::new(move || { + let activations = Arc::clone(&activations); + Box::pin(async move { + activations.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + }), + ) + .unwrap(); + } + + registry.activate_targets().await.unwrap(); + registry.activate_targets().await.unwrap(); + + assert_eq!(policy_calls.load(Ordering::SeqCst), 2); + assert_eq!(activations.load(Ordering::SeqCst), 2); + } + + #[test] + fn target_kind_deserialization_rejects_non_namespaced_values() { + assert!(serde_json::from_value::(json!("otlp_trace")).is_err()); + assert_eq!( + serde_json::from_value::(json!("example.otlp.trace")) + .unwrap() + .as_str(), + "example.otlp.trace" + ); + + let registry = ExportActivationRegistry::default(); + registry + .register_target( + "example:duplicate".into(), + registration("duplicate", None, true), + Arc::new(|| Box::pin(async { Ok(()) })), + ) + .unwrap(); + assert!( + registry + .register_target( + "example:duplicate".into(), + registration("duplicate", None, true), + Arc::new(|| Box::pin(async { Ok(()) })), + ) + .is_err() + ); + } +} diff --git a/crates/core/src/api/mod.rs b/crates/core/src/api/mod.rs index 21b5b83b3..8869eaa05 100644 --- a/crates/core/src/api/mod.rs +++ b/crates/core/src/api/mod.rs @@ -5,6 +5,8 @@ /// Lifecycle event types and builder-backed event constructors. pub mod event; +/// Activation-time policy types for Relay-managed remote exporters. +pub mod export_activation; /// LLM lifecycle helpers and managed execution entry points. pub mod llm; /// Plugin-neutral evidence recording for managed LLM calls. diff --git a/crates/core/src/observability/plugin_component.rs b/crates/core/src/observability/plugin_component.rs index f1496919c..64b3e73f6 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -36,6 +36,11 @@ use serde_json::{Map, Value as Json}; use uuid::Uuid; use crate::api::event::{Event, LogSeverity, ScopeCategory, ValidatedMetricMeasurement}; +use crate::api::export_activation::{ + ExportActivationDecision, ExportActivationPolicyConfig, ExportActivationPolicyRegistry, + ExportActivationRequest, ExportActivationTargetKind, MAX_EXPORT_ACTIVATION_TIMEOUT_MILLIS, + MIN_EXPORT_ACTIVATION_TIMEOUT_MILLIS, +}; use crate::api::runtime::{EventSubscriberFn, current_scope_stack, global_context}; use crate::api::scope::ScopeType; use crate::api::subscriber::{ @@ -223,6 +228,9 @@ pub struct OpenTelemetrySignalEndpointConfig { /// OTLP request timeout in milliseconds. #[serde(default = "default_timeout_millis")] pub timeout_millis: u64, + /// Optional activation-time policy for this destination. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub activation_policy: Option, } /// OTLP log pipeline settings. @@ -363,6 +371,9 @@ pub struct OpenTelemetryEndpointConfig { /// Maximum delay before exporting a non-full batch, in milliseconds. #[serde(default, skip_serializing_if = "Option::is_none")] pub scheduled_delay_millis: Option, + /// Optional activation-time policy for this destination. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub activation_policy: Option, } /// Multi-sink ATOF JSONL exporter config. @@ -406,6 +417,9 @@ pub struct AtofFileSinkSectionConfig { #[serde(default = "default_atof_mode")] #[cfg_attr(feature = "schema", schemars(schema_with = "atof_mode_schema"))] pub mode: String, + /// Optional activation-time policy for this local file sink. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub activation_policy: Option, } /// Stream sink settings for the ATOF plugin section. @@ -436,6 +450,9 @@ pub struct AtofStreamSinkSectionConfig { /// Optional stable name used by other components to reference this endpoint. #[serde(default, skip_serializing_if = "Option::is_none")] pub name: Option, + /// Optional activation-time policy for this remote sink. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub activation_policy: Option, } /// Per-trajectory ATIF exporter config. @@ -478,6 +495,10 @@ pub struct AtifSectionConfig { /// [`storage`]: Self::storage #[serde(default = "default_atif_filename_template")] pub filename_template: String, + /// Optional activation-time policy for the implicit local file destination. + /// This also controls local fallback when remote storage writes fail. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub local_activation_policy: Option, /// Optional list of remote storage destinations. When non-empty, completed /// trajectories are uploaded to every configured backend instead of being /// written locally; the local file write at [`output_directory`] is @@ -501,6 +522,7 @@ impl Default for AtifSectionConfig { extra: None, output_directory: None, filename_template: default_atif_filename_template(), + local_activation_policy: None, storage: Vec::new(), } } @@ -573,6 +595,9 @@ pub struct S3StorageConfig { /// Allow plain HTTP endpoints. When unset, `AWS_ALLOW_HTTP` is used. #[serde(default, skip_serializing_if = "Option::is_none")] pub allow_http: Option, + /// Optional activation-time policy for this remote storage target. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub activation_policy: Option, } /// HTTP endpoint settings for ATIF trajectory upload. @@ -595,6 +620,9 @@ pub struct HttpStorageConfig { /// Request timeout in milliseconds. #[serde(default = "default_timeout_millis")] pub timeout_millis: u64, + /// Optional activation-time policy for this remote storage target. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub activation_policy: Option, } crate::editor_config! { @@ -767,6 +795,7 @@ impl EditorConfig for OpenTelemetryEndpointConfig { &[], false, ), + otel_editor_field("activation_policy", EditorFieldKind::Json, &[], true), ], }; &SCHEMA @@ -797,6 +826,7 @@ impl EditorConfig for OpenTelemetrySignalEndpointConfig { &[], false, ), + otel_editor_field("activation_policy", EditorFieldKind::Json, &[], true), ], }; &SCHEMA @@ -858,6 +888,7 @@ crate::editor_config! { output_directory => { label: "output_directory", kind: String, optional: true }, filename => { label: "filename", kind: String, optional: true }, mode => { label: "mode", kind: Enum, values: ["append", "overwrite"] }, + activation_policy => { label: "activation_policy", kind: Json, optional: true }, } } @@ -870,6 +901,7 @@ crate::editor_config! { timeout_millis => { label: "timeout_millis", kind: Integer }, field_name_policy => { label: "field_name_policy", kind: Enum, values: ["preserve", "replace_dots"] }, name => { label: "name", kind: String, optional: true }, + activation_policy => { label: "activation_policy", kind: Json, optional: true }, } } @@ -883,6 +915,7 @@ crate::editor_config! { extra => { label: "extra", kind: Json, optional: true }, output_directory => { label: "output_directory", kind: String, optional: true }, filename_template => { label: "filename_template", kind: String }, + local_activation_policy => { label: "local_activation_policy", kind: Json, optional: true }, storage => { label: "storage", kind: Json, optional: true }, } } @@ -962,7 +995,7 @@ impl Plugin for ObservabilityPlugin { let plugin_config = plugin_config.clone(); Box::pin(async move { let config = parse_observability_config(&plugin_config)?; - register_observability(config, ctx) + register_observability(config, ctx).await }) } } @@ -1064,7 +1097,7 @@ fn string_enum_schema( schema.into() } -fn register_observability( +async fn register_observability( config: ObservabilityConfig, ctx: &mut PluginRegistrationContext, ) -> PluginResult<()> { @@ -1087,13 +1120,13 @@ fn register_observability( } register_full_payload_policy(config.enable_full_payloads, ctx)?; if let Some(atof) = config.atof.filter(|section| section.enabled) { - register_atof_exporter(atof, ctx)?; + register_atof_exporter(atof, ctx).await?; } if let Some(atif) = config.atif.filter(|section| section.enabled) { - register_atif_dispatcher(atif, ctx)?; + register_atif_dispatcher(atif, ctx).await?; } if let Some(otel) = config.opentelemetry.filter(|section| section.enabled) { - register_opentelemetry(otel, ctx)?; + register_opentelemetry(otel, ctx).await?; } Ok(()) } @@ -1121,14 +1154,42 @@ fn register_full_payload_policy( Ok(()) } -fn register_atof_exporter( +async fn register_atof_exporter( section: AtofSectionConfig, ctx: &mut PluginRegistrationContext, ) -> PluginResult<()> { - let exporters = section - .sinks + let export_activation_policies = ctx.export_activation_policies(); + let mut allowed_sinks = Vec::with_capacity(section.sinks.len()); + for (index, sink) in section.sinks.into_iter().enumerate() { + let allowed = match &sink { + AtofSinkSectionConfig::File(file) => { + export_target_allowed( + &export_activation_policies, + file.activation_policy.as_ref(), + ExportActivationTargetKind::ATOF_FILE, + &format!("atof.sinks[{index}]"), + ) + .await + } + AtofSinkSectionConfig::Stream(stream) => { + export_target_allowed( + &export_activation_policies, + stream.activation_policy.as_ref(), + ExportActivationTargetKind::AtofStream, + &format!("atof.sinks[{index}]"), + ) + .await + } + }; + if allowed { + allowed_sinks.push((index, sink)); + } + } + if allowed_sinks.is_empty() { + return Ok(()); + } + let exporters = allowed_sinks .into_iter() - .enumerate() .map(|(index, sink)| { let config = CoreAtofExporterConfig { sink: build_atof_sink_config(index, sink)?, @@ -1214,20 +1275,70 @@ fn build_atof_sink_config( type AtifStorageList = Arc>>; -fn register_atif_dispatcher( - section: AtifSectionConfig, +async fn register_atif_dispatcher( + mut section: AtifSectionConfig, ctx: &mut PluginRegistrationContext, ) -> PluginResult<()> { + let export_activation_policies = ctx.export_activation_policies(); validate_atif_filename_template(§ion.filename_template) .map_err(PluginError::InvalidConfig)?; - let mut storage_vec = Vec::with_capacity(section.storage.len()); - for (index, entry) in section.storage.iter().enumerate() { + let had_remote_storage = !section.storage.is_empty(); + let local_allowed = export_target_allowed( + &export_activation_policies, + section.local_activation_policy.as_ref(), + ExportActivationTargetKind::ATIF_FILE, + "atif.local", + ) + .await; + let mut allowed_storage = Vec::with_capacity(section.storage.len()); + for (index, entry) in section.storage.into_iter().enumerate() { + let (policy, kind) = match &entry { + AtifStorageConfig::Http(config) => ( + config.activation_policy.as_ref(), + ExportActivationTargetKind::AtifHttp, + ), + AtifStorageConfig::S3(config) => ( + config.activation_policy.as_ref(), + ExportActivationTargetKind::AtifS3, + ), + }; + if export_target_allowed( + &export_activation_policies, + policy, + kind, + &format!("atif.storage[{index}]"), + ) + .await + { + allowed_storage.push((index, entry)); + } + } + if (had_remote_storage && allowed_storage.is_empty()) || (!had_remote_storage && !local_allowed) + { + return Ok(()); + } + + let remote_storage_indices = allowed_storage + .iter() + .map(|(index, _)| *index) + .collect::>(); + let mut storage_vec = Vec::with_capacity(allowed_storage.len()); + for (index, entry) in &allowed_storage { + let index = *index; storage_vec.push(build_atif_storage(index, entry)?); } + section.storage = allowed_storage + .into_iter() + .map(|(_, entry)| entry) + .collect(); let storage: AtifStorageList = Arc::new(storage_vec); - let manager = Arc::new(Mutex::new(AtifDispatcher::new(section))); + let manager = Arc::new(Mutex::new(AtifDispatcher::with_remote_storage_indices( + section, + remote_storage_indices, + local_allowed, + ))); let dispatcher = atif_dispatcher_subscriber( Arc::clone(&manager), ctx.qualify_name("atif-"), @@ -1364,10 +1475,11 @@ fn build_atif_storage( )) } -fn register_opentelemetry( +async fn register_opentelemetry( section: OpenTelemetrySectionConfig, ctx: &mut PluginRegistrationContext, ) -> PluginResult<()> { + let export_activation_policies = ctx.export_activation_policies(); let OpenTelemetrySectionConfig { endpoints, logs, @@ -1391,10 +1503,49 @@ fn register_opentelemetry( .as_ref() .map(|section| resolve_signal_endpoints("metrics", section.endpoints.as_ref(), &endpoints)) .transpose()?; - let trace_subscribers = build_opentelemetry_subscribers(endpoints)?; - let log_subscribers = match (logs, log_endpoints) { - (Some(section), Some(endpoints)) => { - match build_opentelemetry_log_subscribers(section, endpoints) { + let mut allowed_traces = Vec::with_capacity(endpoints.len()); + for (index, endpoint) in endpoints.into_iter().enumerate() { + if export_target_allowed( + &export_activation_policies, + endpoint.activation_policy.as_ref(), + ExportActivationTargetKind::OtlpTrace, + &format!("opentelemetry.traces[{index}]"), + ) + .await + { + allowed_traces.push(endpoint); + } + } + let mut allowed_logs = Vec::new(); + for (index, endpoint) in log_endpoints.unwrap_or_default().into_iter().enumerate() { + if export_target_allowed( + &export_activation_policies, + endpoint.activation_policy.as_ref(), + ExportActivationTargetKind::OtlpLog, + &format!("opentelemetry.logs.endpoints[{index}]"), + ) + .await + { + allowed_logs.push(endpoint); + } + } + let mut allowed_metrics = Vec::new(); + for (index, endpoint) in metric_endpoints.unwrap_or_default().into_iter().enumerate() { + if export_target_allowed( + &export_activation_policies, + endpoint.activation_policy.as_ref(), + ExportActivationTargetKind::OtlpMetric, + &format!("opentelemetry.metrics.endpoints[{index}]"), + ) + .await + { + allowed_metrics.push(endpoint); + } + } + let trace_subscribers = build_opentelemetry_subscribers(allowed_traces)?; + let log_subscribers = match logs { + Some(section) if !allowed_logs.is_empty() => { + match build_opentelemetry_log_subscribers(section, allowed_logs) { Ok(subscribers) => subscribers, Err(error) => { let _ = shutdown_opentelemetry_providers(&trace_subscribers); @@ -1404,9 +1555,9 @@ fn register_opentelemetry( } _ => Vec::new(), }; - let metric_subscribers = match (metrics, metric_endpoints) { - (Some(section), Some(endpoints)) => { - match build_opentelemetry_metric_subscribers(section, endpoints) { + let metric_subscribers = match metrics { + Some(section) if !allowed_metrics.is_empty() => { + match build_opentelemetry_metric_subscribers(section, allowed_metrics) { Ok(subscribers) => subscribers, Err(error) => { let _ = shutdown_opentelemetry_providers(&trace_subscribers); @@ -1419,6 +1570,9 @@ fn register_opentelemetry( } _ => Vec::new(), }; + if trace_subscribers.is_empty() && log_subscribers.is_empty() && metric_subscribers.is_empty() { + return Ok(()); + } for (signal, count) in [ ("traces", trace_subscribers.len()), ("logs", log_subscribers.len()), @@ -1751,6 +1905,7 @@ fn derive_signal_endpoint( service_version: trace.service_version.clone(), instrumentation_scope: trace.instrumentation_scope.clone(), timeout_millis: trace.timeout_millis, + activation_policy: trace.activation_policy.clone(), }) } @@ -2124,6 +2279,8 @@ fn shutdown_opentelemetry_providers( struct AtifDispatcher { config: AtifSectionConfig, + remote_storage_indices: Vec, + local_allowed: bool, agents: HashMap, scope_owners: HashMap, scope_subscribers: HashMap, @@ -2224,9 +2381,21 @@ enum SinkLabel { } impl AtifDispatcher { + #[cfg(test)] fn new(config: AtifSectionConfig) -> Self { + let remote_storage_indices = (0..config.storage.len()).collect(); + Self::with_remote_storage_indices(config, remote_storage_indices, true) + } + + fn with_remote_storage_indices( + config: AtifSectionConfig, + remote_storage_indices: Vec, + local_allowed: bool, + ) -> Self { Self { config, + remote_storage_indices, + local_allowed, agents: HashMap::new(), scope_owners: HashMap::new(), scope_subscribers: HashMap::new(), @@ -2514,14 +2683,19 @@ impl AtifDispatcher { .clone() .unwrap_or_else(default_output_directory); let path = directory.join(&filename); - Ok((filename, Some(path))) + Ok((filename, self.local_allowed.then_some(path))) } fn sink_targets(&self) -> Vec { if self.config.storage.is_empty() { - vec![SinkLabel::Local] + self.local_allowed + .then_some(SinkLabel::Local) + .into_iter() + .collect() } else { - (0..self.config.storage.len()) + self.remote_storage_indices + .iter() + .copied() .map(SinkLabel::Remote) .collect() } @@ -2842,6 +3016,7 @@ fn write_atif( storage: &[Arc], targets: &[SinkLabel], ) -> Vec<(SinkLabel, std::io::Result<()>)> { + let mut remote_position = 0; let mut results = targets .iter() .map(|label| { @@ -2852,7 +3027,11 @@ fn write_atif( "ATIF local destination has no output path", )), }, - SinkLabel::Remote(index) => write_atif_remote(storage, *index, write), + SinkLabel::Remote(index) => { + let result = write_atif_remote(storage, remote_position, *index, write); + remote_position += 1; + result + } }; (label.clone(), result) }) @@ -2862,6 +3041,7 @@ fn write_atif( .iter() .all(|label| matches!(label, SinkLabel::Remote(_))) && results.iter().all(|(_, result)| result.is_err()) + && write.local_path.is_some() { let fallback = match &write.local_path { Some(path) => write_atif_local(path, &write.payload), @@ -2884,19 +3064,21 @@ fn write_atif_local(path: &PathBuf, payload: &[u8]) -> std::io::Result<()> { #[cfg(feature = "object-store")] fn write_atif_remote( storage: &[Arc], - index: usize, + storage_position: usize, + config_index: usize, write: &PendingAtifWrite, ) -> std::io::Result<()> { - let sink = storage - .get(index) - .ok_or_else(|| std::io::Error::other(format!("ATIF storage[{index}] is not registered")))?; + let sink = storage.get(storage_position).ok_or_else(|| { + std::io::Error::other(format!("ATIF storage[{config_index}] is not registered")) + })?; sink.put(&write.filename, &write.session_id, &write.payload) } #[cfg(not(feature = "object-store"))] fn write_atif_remote( _storage: &[Arc], - _index: usize, + _storage_position: usize, + _config_index: usize, _write: &PendingAtifWrite, ) -> std::io::Result<()> { Err(std::io::Error::other( @@ -3182,6 +3364,7 @@ fn validate_observability_section_fields( "extra", "output_directory", "filename_template", + "local_activation_policy", "storage", ], ); @@ -3263,6 +3446,7 @@ fn validate_opentelemetry_signal_fields( "service_version", "instrumentation_scope", "timeout_millis", + "activation_policy", ]; let section_fields = match signal { "logs" => &[ @@ -3344,6 +3528,7 @@ fn validate_opentelemetry_endpoint_fields( "max_queue_size", "max_export_batch_size", "scheduled_delay_millis", + "activation_policy", ]; const REMOVED: &[&str] = &["semantic_selector", "capture_content"]; let Some(endpoints) = opentelemetry @@ -3566,6 +3751,12 @@ fn validate_opentelemetry_section( } validate_opentelemetry_batch_config(diagnostics, policy, index, endpoint); validate_opentelemetry_headers(diagnostics, policy, index, endpoint); + validate_export_activation_policy( + diagnostics, + policy, + &format!("opentelemetry.endpoints[{index}].activation_policy"), + endpoint.activation_policy.as_ref(), + ); } for error in opentelemetry_destination_collision_errors(§ion.endpoints) { diagnostics.push(ConfigDiagnostic { @@ -3743,6 +3934,12 @@ fn validate_opentelemetry_signal_endpoint_values( "headers", endpoint.headers.keys(), ); + validate_export_activation_policy( + diagnostics, + policy, + &format!("opentelemetry.{signal}.endpoints[{index}].activation_policy"), + endpoint.activation_policy.as_ref(), + ); validate_case_insensitive_signal_header_duplicates( diagnostics, policy, @@ -4245,6 +4442,12 @@ fn validate_atof_sink<'a>( ) { match sink { AtofSinkSectionConfig::File(file) => { + validate_export_activation_policy( + diagnostics, + policy, + &format!("atof.sinks[{index}].activation_policy"), + file.activation_policy.as_ref(), + ); if AtofExporterMode::parse(&file.mode).is_none() { push_policy_diag( diagnostics, @@ -4303,6 +4506,12 @@ fn validate_atof_stream_sink_values( index: usize, endpoint: &AtofStreamSinkSectionConfig, ) { + validate_export_activation_policy( + diagnostics, + policy, + &format!("atof.sinks[{index}].activation_policy"), + endpoint.activation_policy.as_ref(), + ); let transport = AtofEndpointTransport::parse(&endpoint.transport); if endpoint.url.trim().is_empty() { push_policy_diag( @@ -4516,6 +4725,12 @@ fn validate_atif_values( policy: &ConfigPolicy, section: &AtifSectionConfig, ) { + validate_export_activation_policy( + diagnostics, + policy, + "atif.local_activation_policy", + section.local_activation_policy.as_ref(), + ); if let Err(message) = validate_atif_filename_template(§ion.filename_template) { push_policy_diag( diagnostics, @@ -4539,6 +4754,12 @@ fn validate_atif_storage_values( ) { match storage { AtifStorageConfig::Http(http) => { + validate_export_activation_policy( + diagnostics, + policy, + &format!("atif.storage[{index}].activation_policy"), + http.activation_policy.as_ref(), + ); validate_atif_http_endpoint( diagnostics, policy, @@ -4580,6 +4801,12 @@ fn validate_atif_storage_values( } } AtifStorageConfig::S3(s3) => { + validate_export_activation_policy( + diagnostics, + policy, + &format!("atif.storage[{index}].activation_policy"), + s3.activation_policy.as_ref(), + ); if s3.bucket.trim().is_empty() { push_policy_diag( diagnostics, @@ -4606,6 +4833,42 @@ fn validate_atif_storage_values( } } +fn validate_export_activation_policy( + diagnostics: &mut Vec, + policy: &ConfigPolicy, + field: &str, + activation_policy: Option<&ExportActivationPolicyConfig>, +) { + let Some(activation_policy) = activation_policy else { + return; + }; + if activation_policy.provider.trim().is_empty() + || activation_policy.provider.trim() != activation_policy.provider + { + push_policy_diag( + diagnostics, + policy.unsupported_value, + "observability.unsupported_value", + Some(OBSERVABILITY_PLUGIN_KIND.to_string()), + Some(format!("{field}.provider")), + "export activation policy provider must be nonblank and have no surrounding whitespace" + .to_string(), + ); + } + if !(MIN_EXPORT_ACTIVATION_TIMEOUT_MILLIS..=MAX_EXPORT_ACTIVATION_TIMEOUT_MILLIS) + .contains(&activation_policy.timeout_millis) + { + push_policy_diag( + diagnostics, + policy.unsupported_value, + "observability.unsupported_value", + Some(OBSERVABILITY_PLUGIN_KIND.to_string()), + Some(format!("{field}.timeout_millis")), + "export activation policy timeout_millis must be between 1000 and 300000".to_string(), + ); + } +} + fn validate_atif_http_header( diagnostics: &mut Vec, policy: &ConfigPolicy, @@ -4797,6 +5060,53 @@ fn observability_registration_error(error: impl std::fmt::Display) -> PluginErro PluginError::RegistrationFailed(error.to_string()) } +async fn export_target_allowed( + export_activation_policies: &ExportActivationPolicyRegistry, + policy: Option<&ExportActivationPolicyConfig>, + target_kind: ExportActivationTargetKind, + field: &str, +) -> bool { + let Some(policy) = policy else { + return true; + }; + let request = ExportActivationRequest { + target_kind: target_kind.clone(), + config: policy.config.clone(), + }; + let outcome = tokio::time::timeout( + export_activation_timeout(policy.timeout_millis), + export_activation_policies.evaluate(&policy.provider, request), + ) + .await; + let (allowed, reason) = match outcome { + Ok(Ok(ExportActivationDecision::Allow)) => (true, "allowed"), + Ok(Ok(ExportActivationDecision::Deny)) => (false, "denied"), + Ok(Err(FlowError::NotFound(_))) => (false, "provider_unavailable"), + Ok(Err(_)) => (false, "provider_error"), + Err(_) => (false, "timeout"), + }; + if !allowed { + log::warn!( + target: "nemo_relay.plugin", + event = "export_activation_policy_denied", + plugin_kind = OBSERVABILITY_PLUGIN_KIND, + provider = policy.provider.as_str(), + target_kind = target_kind.as_str(), + field, + reason; + "Exporter target suppressed by activation policy" + ); + } + allowed +} + +fn export_activation_timeout(timeout_millis: u64) -> Duration { + Duration::from_millis(timeout_millis.clamp( + MIN_EXPORT_ACTIVATION_TIMEOUT_MILLIS, + MAX_EXPORT_ACTIVATION_TIMEOUT_MILLIS, + )) +} + fn default_observability_config_version() -> u32 { 4 } diff --git a/crates/core/src/plugin.rs b/crates/core/src/plugin.rs index 4c37f2c16..3006cf49f 100644 --- a/crates/core/src/plugin.rs +++ b/crates/core/src/plugin.rs @@ -22,6 +22,10 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value as Json}; use thiserror::Error; +use crate::api::export_activation::{ + ExportActivationPolicyFn, ExportActivationPolicyRegistry, ExportTargetActivationFn, + ExportTargetRegistration, +}; use crate::api::registry::{ deregister_event_metadata_injector, deregister_llm_conditional_execution_guardrail, deregister_llm_execution_intercept, deregister_llm_request_intercept, @@ -401,10 +405,22 @@ impl PluginRegistration { /// Each `register_*` call both installs the middleware/subscriber into the /// NeMo Relay runtime and records the inverse deregistration closure so the host /// can roll back partial setup on failure. -#[derive(Default)] pub struct PluginRegistrationContext { registrations: Vec, namespace: Option, + provider_id: Option, + export_activation_policies: Arc, +} + +impl Default for PluginRegistrationContext { + fn default() -> Self { + Self { + registrations: Vec::new(), + namespace: None, + provider_id: None, + export_activation_policies: Arc::new(ExportActivationPolicyRegistry::default()), + } + } } impl PluginRegistrationContext { @@ -418,9 +434,46 @@ impl PluginRegistrationContext { Self { registrations: vec![], namespace: Some(namespace.into()), + provider_id: None, + export_activation_policies: Arc::new(ExportActivationPolicyRegistry::default()), + } + } + + pub(crate) fn with_namespace_and_export_activation_policies( + namespace: impl Into, + provider_id: impl Into, + export_activation_policies: Arc, + ) -> Self { + Self { + registrations: Vec::new(), + namespace: Some(namespace.into()), + provider_id: Some(provider_id.into()), + export_activation_policies, } } + #[cfg(test)] + pub(crate) fn with_export_activation_policies( + export_activation_policies: Arc, + ) -> Self { + Self { + registrations: Vec::new(), + namespace: None, + provider_id: None, + export_activation_policies, + } + } + + #[doc(hidden)] + pub fn export_activation_policies(&self) -> Arc { + Arc::clone(&self.export_activation_policies) + } + + #[doc(hidden)] + pub fn export_activation_provider_id(&self) -> Option<&str> { + self.provider_id.as_deref() + } + /// Returns the runtime-qualified name for a plugin-local registration. /// /// Plugin handlers should pass stable component-local names such as @@ -433,6 +486,84 @@ impl PluginRegistrationContext { } } + /// Registers the single export-activation policy owned by `provider`. + pub(crate) fn register_export_activation_policy_for_provider( + &mut self, + provider: &str, + callback: ExportActivationPolicyFn, + ) -> Result<()> { + self.export_activation_policies + .register(provider, callback) + .map_err(|error| { + PluginError::RegistrationFailed(format!("export activation policy: {error}")) + })?; + + let provider = provider.to_string(); + let export_activation_policies = Arc::clone(&self.export_activation_policies); + self.registrations.push(PluginRegistration::new( + "export_activation_policy", + provider.clone(), + Box::new(move || { + export_activation_policies + .deregister(&provider) + .map(|_| ()) + .map_err(|error| { + PluginError::RegistrationFailed(format!( + "export activation policy deregistration failed: {error}" + )) + }) + }), + )); + Ok(()) + } + + /// Registers the activation policy owned by this component's plugin kind. + pub fn register_export_activation_policy( + &mut self, + callback: ExportActivationPolicyFn, + ) -> Result<()> { + let provider = self.provider_id.clone().ok_or_else(|| { + PluginError::RegistrationFailed( + "component export activation policies require a host-owned plugin context".into(), + ) + })?; + self.register_export_activation_policy_for_provider(&provider, callback) + } + + /// Registers one deferred local or remote export target. + pub fn register_export_target( + &mut self, + registration: ExportTargetRegistration, + activate: F, + ) -> Result<()> + where + F: Fn() -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let qualified_id = self.qualify_name(&format!("export-target:{}", registration.id)); + let activation: ExportTargetActivationFn = Arc::new(move || Box::pin(activate())); + self.export_activation_policies + .register_target(qualified_id.clone(), registration, activation) + .map_err(|error| PluginError::RegistrationFailed(format!("export target: {error}")))?; + + let export_activation_policies = Arc::clone(&self.export_activation_policies); + self.registrations.push(PluginRegistration::new( + "export_target", + qualified_id.clone(), + Box::new(move || { + export_activation_policies + .deregister_target(&qualified_id) + .map(|_| ()) + .map_err(|error| { + PluginError::RegistrationFailed(format!( + "export target deregistration failed: {error}" + )) + }) + }), + )); + Ok(()) + } + /// Registers an event subscriber and records its rollback closure. pub fn register_subscriber(&mut self, name: &str, callback: EventSubscriberFn) -> Result<()> { let qualified_name = self.qualify_name(name); @@ -1486,10 +1617,12 @@ async fn initialize_plugins_with_diagnostics( run_owned_plugin_mutation("plugin initialization", move || async move { let lease = LegacyPluginMutationLease::acquire()?; let rollback_failures = Arc::new(Mutex::new(Vec::new())); + let export_activation_policies = Arc::new(ExportActivationPolicyRegistry::default()); let initialization = tokio::spawn(initialize_plugins_exact_inner( config, Some(Arc::clone(&rollback_failures)), diagnostics, + export_activation_policies, )) .await .map_err(|error| { @@ -1607,15 +1740,23 @@ pub(crate) async fn initialize_plugins_exact_for_host( owner_id: u64, rollback_failures: Arc>>, diagnostics: Vec, + export_activation_policies: Arc, ) -> Result { verify_plugin_host_owner(owner_id)?; - initialize_plugins_exact_inner(config, Some(rollback_failures), diagnostics).await + initialize_plugins_exact_inner( + config, + Some(rollback_failures), + diagnostics, + export_activation_policies, + ) + .await } async fn initialize_plugins_exact_inner( config: PluginConfig, rollback_failures: Option>>>, diagnostics: Vec, + export_activation_policies: Arc, ) -> Result { let enabled_component_count = config .components @@ -1654,6 +1795,7 @@ async fn initialize_plugins_exact_inner( previous_state, rollback_failures, enabled_component_count, + export_activation_policies, ) .await } @@ -1663,6 +1805,7 @@ async fn initialize_plugins_exact_inner( report, rollback_failures, enabled_component_count, + export_activation_policies, ) .await } @@ -1674,10 +1817,20 @@ async fn activate_initial_plugin_configuration( report: ConfigReport, rollback_failures: Option>>>, enabled_component_count: usize, + export_activation_policies: Arc, ) -> Result { - let registrations = - initialize_plugin_components_catching_panics(config.clone(), rollback_failures).await?; - store_active_plugin_configuration(config, report.clone(), registrations)?; + let registrations = initialize_plugin_components_catching_panics( + config.clone(), + rollback_failures, + Arc::clone(&export_activation_policies), + ) + .await?; + store_active_plugin_configuration_with_export_activation_policies( + config, + report.clone(), + registrations, + export_activation_policies, + )?; log::info!( target: "nemo_relay.plugin", event = "plugin_configuration_activated", @@ -1693,6 +1846,7 @@ async fn replace_plugin_configuration( mut previous_state: ActivePluginConfiguration, rollback_failures: Option>>>, enabled_component_count: usize, + export_activation_policies: Arc, ) -> Result { install_previous_configuration_for_teardown(&previous_state)?; let teardown = rollback_registrations_checked(&mut previous_state.registrations); @@ -1710,6 +1864,7 @@ async fn replace_plugin_configuration( previous_state, rollback_failures, enabled_component_count, + export_activation_policies, ) .await } @@ -1725,6 +1880,7 @@ fn install_previous_configuration_for_teardown( report: previous_state.report.clone(), runtime_diagnostics: previous_state.runtime_diagnostics.clone(), registrations: Vec::new(), + export_activation_policies: Arc::clone(&previous_state.export_activation_policies), }); Ok(()) } @@ -1760,12 +1916,22 @@ async fn activate_replacement_or_restore( previous_state: ActivePluginConfiguration, rollback_failures: Option>>>, enabled_component_count: usize, + export_activation_policies: Arc, ) -> Result { - match initialize_plugin_components_catching_panics(config.clone(), rollback_failures.clone()) - .await + match initialize_plugin_components_catching_panics( + config.clone(), + rollback_failures.clone(), + Arc::clone(&export_activation_policies), + ) + .await { Ok(registrations) => { - store_active_plugin_configuration(config, report.clone(), registrations)?; + store_active_plugin_configuration_with_export_activation_policies( + config, + report.clone(), + registrations, + export_activation_policies, + )?; log::info!( target: "nemo_relay.plugin", event = "plugin_configuration_replaced", @@ -1785,9 +1951,11 @@ async fn restore_previous_plugin_configuration( rollback_failures: Option>>>, err: PluginError, ) -> Result { + let export_activation_policies = Arc::new(ExportActivationPolicyRegistry::default()); match initialize_plugin_components_catching_panics( previous_state.config.clone(), rollback_failures, + Arc::clone(&export_activation_policies), ) .await { @@ -1797,6 +1965,7 @@ async fn restore_previous_plugin_configuration( previous_state.report, previous_state.runtime_diagnostics, registrations, + export_activation_policies, )?; log::warn!( target: "nemo_relay.plugin", @@ -1823,14 +1992,17 @@ async fn restore_previous_plugin_configuration( async fn initialize_plugin_components_catching_panics( config: PluginConfig, rollback_failures: Option>>>, + export_activation_policies: Arc, ) -> Result> { - tokio::spawn(async move { initialize_plugin_components(&config, rollback_failures).await }) - .await - .map_err(|error| { - PluginError::Internal(format!( - "plugin component initialization task failed: {error}" - )) - })? + tokio::spawn(async move { + initialize_plugin_components(&config, rollback_failures, export_activation_policies).await + }) + .await + .map_err(|error| { + PluginError::Internal(format!( + "plugin component initialization task failed: {error}" + )) + })? } /// Validates and activates `config` layered on top of the discovered @@ -2606,21 +2778,30 @@ struct ActivePluginConfiguration { report: ConfigReport, runtime_diagnostics: BTreeMap, registrations: Vec, + export_activation_policies: Arc, } async fn initialize_plugin_components( config: &PluginConfig, rollback_failures: Option>>>, + export_activation_policies: Arc, ) -> Result> { ensure_builtin_plugins_registered()?; let totals = plugin_component_totals(config); let mut ordinals: HashMap<&str, usize> = HashMap::new(); let mut registrations = PendingPluginRegistrations::new(rollback_failures.clone()); - for component in config + // Observability is the built-in consumer of export policies. Register all + // other components first so ordinary and dynamic policy providers are + // complete before observability evaluates and constructs its targets. + let enabled_components = config .components .iter() - .filter(|component| component.enabled) + .filter(|component| component.enabled); + for component in enabled_components + .clone() + .filter(|component| component.kind != "observability") + .chain(enabled_components.filter(|component| component.kind == "observability")) { let Some(plugin) = lookup_registered_plugin(&component.kind) else { return Err(PluginError::NotFound(format!( @@ -2639,14 +2820,25 @@ async fn initialize_plugin_components( totals.get(component.kind.as_str()).copied().unwrap_or(1), ); - let mut pending = - PendingPluginRegistrationContext::new(namespace, rollback_failures.clone()); + let mut pending = PendingPluginRegistrationContext::new( + namespace, + component.kind.clone(), + rollback_failures.clone(), + Arc::clone(&export_activation_policies), + ); plugin .register(&component.config, &mut pending.context) .await?; registrations.extend(pending.take()); } + export_activation_policies + .activate_targets() + .await + .map_err(|error| { + PluginError::RegistrationFailed(format!("export target activation failed: {error}")) + })?; + Ok(registrations.take()) } @@ -2687,9 +2879,18 @@ struct PendingPluginRegistrationContext { } impl PendingPluginRegistrationContext { - fn new(namespace: String, rollback_failures: Option>>>) -> Self { + fn new( + namespace: String, + provider_id: String, + rollback_failures: Option>>>, + export_activation_policies: Arc, + ) -> Self { Self { - context: PluginRegistrationContext::with_namespace(namespace), + context: PluginRegistrationContext::with_namespace_and_export_activation_policies( + namespace, + provider_id, + export_activation_policies, + ), rollback_failures, } } @@ -2722,16 +2923,32 @@ fn record_rollback_failures( } } +#[cfg(test)] fn store_active_plugin_configuration( config: PluginConfig, report: ConfigReport, registrations: Vec, +) -> Result<()> { + store_active_plugin_configuration_with_export_activation_policies( + config, + report, + registrations, + Arc::new(ExportActivationPolicyRegistry::default()), + ) +} + +fn store_active_plugin_configuration_with_export_activation_policies( + config: PluginConfig, + report: ConfigReport, + registrations: Vec, + export_activation_policies: Arc, ) -> Result<()> { store_active_plugin_configuration_with_runtime_diagnostics( config, report, BTreeMap::new(), registrations, + export_activation_policies, ) } @@ -2740,6 +2957,7 @@ fn store_active_plugin_configuration_with_runtime_diagnostics( report: ConfigReport, runtime_diagnostics: BTreeMap, registrations: Vec, + export_activation_policies: Arc, ) -> Result<()> { let mut guard = ACTIVE_PLUGIN_CONFIGURATION.lock().map_err(|err| { PluginError::Internal(format!("active plugin configuration lock poisoned: {err}")) @@ -2749,6 +2967,7 @@ fn store_active_plugin_configuration_with_runtime_diagnostics( report, runtime_diagnostics, registrations, + export_activation_policies, }); if let Ok(mut guard) = LAST_FAILED_RUNTIME_DIAGNOSTICS_REPORT.lock() { *guard = None; diff --git a/crates/core/src/plugin/dynamic.rs b/crates/core/src/plugin/dynamic.rs index 6fbedc526..9add23fa7 100644 --- a/crates/core/src/plugin/dynamic.rs +++ b/crates/core/src/plugin/dynamic.rs @@ -244,6 +244,8 @@ pub enum DynamicPluginCapability { PluginWorker, /// Typed configuration schema contribution capability. ConfigSchema, + /// Activation-time policy for local or remote export targets. + ExportActivationPolicy, } /// Host policy startup classification for a plugin. diff --git a/crates/core/src/plugin/dynamic/host.rs b/crates/core/src/plugin/dynamic/host.rs index 49eb57579..1dae73c1b 100644 --- a/crates/core/src/plugin/dynamic/host.rs +++ b/crates/core/src/plugin/dynamic/host.rs @@ -15,6 +15,7 @@ use std::sync::{Arc, Mutex}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value as Json}; +use crate::api::export_activation::ExportActivationPolicyRegistry; use crate::plugin::{ ConfigReport, PluginComponentSpec, PluginConfig, PluginHostLease, Result, acquire_plugin_host_lease, clear_plugin_configuration_for_host, @@ -23,7 +24,8 @@ use crate::plugin::{ }; use super::{ - DynamicPluginKind, DynamicPluginTeardownOutcome, NativePluginActivation, NativePluginLoadSpec, + DynamicPluginCapability, DynamicPluginKind, DynamicPluginManifest, + DynamicPluginTeardownOutcome, NativePluginActivation, NativePluginLoadSpec, load_native_plugins, }; @@ -56,6 +58,7 @@ pub struct DynamicPluginActivationSpec { #[must_use = "dropping the activation clears and unloads its dynamic plugins"] pub struct PluginHostActivation { active: bool, + export_activation_policies: Option>, native: Option, #[cfg(feature = "worker-grpc")] worker: Option, @@ -177,22 +180,37 @@ impl PluginHostActivation { .transpose()? }; - config.components.extend( - dynamic_plugins - .into_iter() - .map(|plugin| PluginComponentSpec { - kind: plugin.plugin_id, - enabled: true, - config: plugin.config, - }), - ); + let mut policy_components = Vec::new(); + let mut regular_components = Vec::new(); + for plugin in dynamic_plugins { + let (manifest, _) = DynamicPluginManifest::load_from_path(&plugin.manifest_ref)?; + let component = PluginComponentSpec { + kind: plugin.plugin_id, + enabled: true, + config: plugin.config, + }; + if manifest + .capabilities + .items + .contains(&DynamicPluginCapability::ExportActivationPolicy) + { + policy_components.push(component); + } else { + regular_components.push(component); + } + } + policy_components.append(&mut config.components); + policy_components.append(&mut regular_components); + config.components = policy_components; let rollback_failures = Arc::new(Mutex::new(Vec::new())); + let export_activation_policies = Arc::new(ExportActivationPolicyRegistry::default()); let owner_id = claim.owner_id(); let initialization = tokio::spawn(initialize_plugins_exact_for_host( config, owner_id, Arc::clone(&rollback_failures), diagnostics, + Arc::clone(&export_activation_policies), )) .await .map_err(|error| { @@ -247,6 +265,7 @@ impl PluginHostActivation { Ok(( Self { active: true, + export_activation_policies: Some(export_activation_policies), native, #[cfg(feature = "worker-grpc")] worker, @@ -295,6 +314,7 @@ impl PluginHostActivation { self.retain_loaded_runtimes(); return Err(retained_runtime_error(errors)); } + self.export_activation_policies.take(); let mut runtime_outcome = DynamicPluginTeardownOutcome::success(); if let Some(native) = &mut self.native { diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 4bb6e1b99..bae5a268b 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -22,6 +22,10 @@ use std::task::{Context, Poll}; use futures_util::FutureExt; use crate::api::event::{DataSchema, Event, EventSanitizeFields, LogSeverity}; +use crate::api::export_activation::{ + ExportActivationDecision, ExportActivationPolicyFn, ExportActivationRequest, + ExportTargetActivationFn, ExportTargetRegistration, +}; use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; use crate::api::runtime::{ EventMetadataInjectorFn, EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, @@ -53,10 +57,11 @@ use chrono::{DateTime, Utc}; use libloading::{Library, Symbol}; use nemo_relay_plugin::{ NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, - NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, - NemoRelayNativeAsyncLlmStreamOpenCb, NemoRelayNativeAsyncLlmStreamPullCb, - NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, - NemoRelayNativeAsyncNextResultCb, NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, + NemoRelayNativeActivationHookKind, NemoRelayNativeAsyncCallbackState, + NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncLlmStreamOpenCb, + NemoRelayNativeAsyncLlmStreamPullCb, NemoRelayNativeAsyncMiddlewareCb, + NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncNextResultCb, + NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativeLlmAsyncStream, @@ -77,7 +82,7 @@ use tokio::runtime::Runtime; use tokio_stream::{Stream, StreamExt}; use super::{ - DynamicPluginKind, DynamicPluginManifest, DynamicPluginManifestLoad, + DynamicPluginCapability, DynamicPluginKind, DynamicPluginManifest, DynamicPluginManifestLoad, DynamicPluginTeardownOutcome, deregister_tracked_registrations_checked, validate_annotated_request_consumer_compatibility, validate_dynamic_plugin_relay_compatibility, }; @@ -282,6 +287,7 @@ struct NativePluginInstance { plugin_kind: String, relay_compat: String, allows_multiple_components: bool, + allows_export_activation_policy: bool, plugin: Mutex, _library: Library, } @@ -442,6 +448,10 @@ fn load_one_native_plugin( plugin_kind, relay_compat, allows_multiple_components: plugin.allows_multiple_components, + allows_export_activation_policy: manifest + .capabilities + .items + .contains(&DynamicPluginCapability::ExportActivationPolicy), plugin: Mutex::new(plugin), _library: library, })) @@ -3504,6 +3514,56 @@ fn wrap_native_async_event_metadata_injector( }) } +fn wrap_native_async_export_activation_policy( + instance: Arc, + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> ExportActivationPolicyFn { + let user_data = make_user_data(instance, user_data, free_fn); + Arc::new(move |request: ExportActivationRequest| { + let user_data = user_data.clone(); + Box::pin(async move { + serde_json::from_value::( + invoke_native_async_callback( + cb, + user_data, + serde_json::to_value(request).map_err(|error| { + FlowError::Internal(format!( + "failed to serialize export activation request: {error}" + )) + })?, + None, + None, + ) + .await?, + ) + .map_err(|error| { + FlowError::Internal(format!( + "invalid native export activation decision: {error}" + )) + }) + }) + }) +} + +fn wrap_native_async_export_target( + instance: Arc, + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> ExportTargetActivationFn { + let user_data = make_user_data(instance, user_data, free_fn); + Arc::new(move || { + let user_data = Arc::clone(&user_data); + Box::pin(async move { + invoke_native_async_callback(cb, user_data, Json::Null, None, None) + .await + .map(|_| ()) + }) + }) +} + fn wrap_native_async_tool_execution( instance: Arc, cb: NemoRelayNativeAsyncMiddlewareCb, @@ -3700,6 +3760,45 @@ unsafe extern "C" fn native_plugin_context_register_async_middleware( Ok(name) => name, Err(status) => return status, }; + if let Ok(hook) = NemoRelayNativeActivationHookKind::try_from(kind) { + let context = unsafe { &mut *host_ctx.ctx }; + let registration = match hook { + NemoRelayNativeActivationHookKind::ExportActivationPolicy => { + if !instance.allows_export_activation_policy { + log::warn!( + target: "nemo_relay.plugin", + event = "native_registration_capability_undeclared", + plugin_id = instance.plugin_kind.as_str(), + capability = "export_activation_policy"; + "Native export activation policy ignored because the manifest capability was not declared" + ); + return NemoRelayStatus::Ok; + } + let (user_data, free_fn) = user_data_guard.transfer(); + context.register_export_activation_policy( + wrap_native_async_export_activation_policy(instance, cb, user_data, free_fn), + ) + } + NemoRelayNativeActivationHookKind::ExportTarget => { + let target = match serde_json::from_str::(&name) { + Ok(target) => target, + Err(error) => { + set_native_last_error(format!( + "invalid native export target registration: {error}" + )); + return NemoRelayStatus::InvalidArg; + } + }; + let (user_data, free_fn) = user_data_guard.transfer(); + let activation = wrap_native_async_export_target(instance, cb, user_data, free_fn); + context.register_export_target(target, move || activation()) + } + }; + return match registration { + Ok(()) => NemoRelayStatus::Ok, + Err(error) => status_from_plugin_error(error), + }; + } let kind = match NemoRelayNativeAsyncMiddlewareKind::try_from(kind) { Ok(kind) => kind, Err(()) => { diff --git a/crates/core/src/plugin/dynamic/worker.rs b/crates/core/src/plugin/dynamic/worker.rs index 4d276db5d..421a464ce 100644 --- a/crates/core/src/plugin/dynamic/worker.rs +++ b/crates/core/src/plugin/dynamic/worker.rs @@ -64,6 +64,10 @@ use tokio_stream::wrappers::UnixListenerStream; use tower::service_fn; use crate::api::event::{DataSchema, Event, EventSanitizeFields, LogSeverity}; +use crate::api::export_activation::{ + ExportActivationDecision, ExportActivationPolicyFn, ExportActivationRequest, + ExportTargetRegistration, +}; use crate::api::llm::{LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, LlmRequest}; use crate::api::runtime::subscriber_dispatcher::{ PublicationBuffer, capture_nested_publication_buffer, with_nested_publication_buffer, @@ -88,7 +92,7 @@ use crate::plugin::{ }; use super::{ - DynamicPluginKind, DynamicPluginManifest, DynamicPluginManifestLoad, + DynamicPluginCapability, DynamicPluginKind, DynamicPluginManifest, DynamicPluginManifestLoad, DynamicPluginTeardownOutcome, WorkerRuntime, deregister_tracked_registrations_checked, validate_annotated_request_consumer_compatibility, validate_dynamic_plugin_relay_compatibility, }; @@ -96,6 +100,8 @@ use super::{ const JSON_SCHEMA: &str = "nemo.relay.Json@1"; const DATA_SCHEMA_SCHEMA: &str = "nemo.relay.DataSchema@1"; const EVENT_SCHEMA: &str = "nemo.relay.Event@1"; +const EXPORT_ACTIVATION_REQUEST_SCHEMA: &str = "nemo.relay.ExportActivationRequest@1"; +const EXPORT_TARGET_REGISTRATION_SCHEMA: &str = "nemo.relay.ExportTargetRegistration@1"; const LLM_REQUEST_SCHEMA: &str = "nemo.relay.LlmRequest@1"; const WORKER_STARTUP_TIMEOUT: Duration = Duration::from_secs(10); const WORKER_RPC_TIMEOUT: Duration = Duration::from_secs(30); @@ -611,11 +617,20 @@ fn load_one_worker_plugin( .map_err(|err| { PluginError::RegistrationFailed(format!("worker registration RPC failed: {err}")) })?; - let register = register.into_inner(); + let mut register = register.into_inner(); if let Some(error) = register.error { return Err(worker_error_to_plugin(error, "worker registration failed")); } validate_registration_plan(&spec.plugin_id, ®ister)?; + filter_unadvertised_export_activation_hooks( + &spec.plugin_id, + manifest + .capabilities + .items + .contains(&DynamicPluginCapability::ExportActivationPolicy), + &handshake.supported_surfaces, + &mut register.registrations, + ); register.registrations }; if registrations.iter().any(|registration| { @@ -1071,6 +1086,11 @@ impl WorkerPluginInstance { ®istration.local_name, registration.priority, )?, + RegistrationSurface::ExportActivationPolicy => self + .install_export_activation_policy_registration(ctx, ®istration.local_name)?, + RegistrationSurface::ExportTarget => { + self.install_export_target_registration(ctx, registration)? + } RegistrationSurface::MarkSanitizeGuardrail | RegistrationSurface::ScopeSanitizeStartGuardrail | RegistrationSurface::ScopeSanitizeEndGuardrail => self @@ -1143,6 +1163,58 @@ impl WorkerPluginInstance { ctx.register_event_metadata_injector(name, priority, callback) } + fn install_export_activation_policy_registration( + &self, + ctx: &mut PluginRegistrationContext, + name: &str, + ) -> crate::plugin::Result<()> { + let callback = Arc::new(self.clone_for_callback()); + let callback_name = name.to_owned(); + let policy: ExportActivationPolicyFn = Arc::new(move |request| { + let callback = Arc::clone(&callback); + let callback_name = callback_name.clone(); + Box::pin(async move { + callback + .invoke_export_activation_policy(&callback_name, request) + .await + }) + }); + ctx.register_export_activation_policy(policy) + } + + fn install_export_target_registration( + &self, + ctx: &mut PluginRegistrationContext, + registration: &Registration, + ) -> crate::plugin::Result<()> { + let envelope = registration.export_target.as_ref().ok_or_else(|| { + PluginError::RegistrationFailed(format!( + "worker export target '{}' omitted registration metadata", + registration.local_name + )) + })?; + if envelope.schema != EXPORT_TARGET_REGISTRATION_SCHEMA { + return Err(PluginError::RegistrationFailed(format!( + "worker export target '{}' used unsupported schema {:?}", + registration.local_name, envelope.schema + ))); + } + let target = + decode_json_envelope::(envelope).map_err(|error| { + PluginError::RegistrationFailed(format!( + "worker export target '{}' metadata is invalid: {error}", + registration.local_name + )) + })?; + let instance = Arc::new(self.clone_for_callback()); + let callback_name = registration.local_name.clone(); + ctx.register_export_target(target, move || { + let instance = Arc::clone(&instance); + let callback_name = callback_name.clone(); + async move { instance.invoke_export_target(&callback_name).await } + }) + } + fn install_event_sanitize_registration( &self, ctx: &mut PluginRegistrationContext, @@ -1545,6 +1617,50 @@ impl WorkerPluginCallback { Ok(additions) } + async fn invoke_export_activation_policy( + &self, + registration_name: &str, + request: ExportActivationRequest, + ) -> FlowResult { + let request = self.base_request( + registration_name, + RegistrationSurface::ExportActivationPolicy, + None, + Some(invoke_request_payload::Payload::ExportActivation( + json_envelope(EXPORT_ACTIVATION_REQUEST_SCHEMA, &request).map_err(|error| { + FlowError::Internal(format!( + "failed to serialize export activation request: {error}" + )) + })?, + )), + ); + let value = json_from_invoke_response(self.invoke_async(request).await?)?; + serde_json::from_value(value).map_err(|error| { + FlowError::Internal(format!( + "worker returned invalid export activation decision: {error}" + )) + }) + } + + async fn invoke_export_target(&self, registration_name: &str) -> FlowResult<()> { + let request = self.base_request( + registration_name, + RegistrationSurface::ExportTarget, + None, + Some(invoke_request_payload::Payload::ExportTarget( + json_envelope("nemo.relay.ExportTargetActivation@1", &Json::Null).map_err( + |error| { + FlowError::Internal(format!( + "failed to serialize export target activation: {error}" + )) + }, + )?, + )), + ); + let _ = json_from_invoke_response(self.invoke_async(request).await?)?; + Ok(()) + } + async fn invoke_event_sanitize( &self, registration_name: &str, @@ -3340,6 +3456,49 @@ fn validate_registration_plan( Ok(()) } +fn filter_unadvertised_export_activation_hooks( + plugin_id: &str, + policy_capability_declared: bool, + supported_surfaces: &[i32], + registrations: &mut Vec, +) { + registrations.retain(|registration| { + let surface = RegistrationSurface::try_from(registration.surface).ok(); + if surface == Some(RegistrationSurface::ExportActivationPolicy) + && !policy_capability_declared + { + log::warn!( + target: "nemo_relay.worker", + event = "worker_registration_capability_undeclared", + plugin_id, + capability = "export_activation_policy"; + "Worker export activation policy ignored because the manifest capability was not declared" + ); + return false; + } + if !matches!( + surface, + Some(RegistrationSurface::ExportActivationPolicy | RegistrationSurface::ExportTarget) + ) || supported_surfaces.contains(®istration.surface) + { + return true; + } + let surface_name = match surface { + Some(RegistrationSurface::ExportActivationPolicy) => "export_activation_policy", + Some(RegistrationSurface::ExportTarget) => "export_target", + _ => unreachable!(), + }; + log::warn!( + target: "nemo_relay.worker", + event = "worker_registration_surface_unadvertised", + plugin_id, + surface = surface_name; + "Worker export activation hook registration ignored because the surface was not advertised" + ); + false + }); +} + fn diagnostics_have_errors(diagnostics: &[ConfigDiagnostic]) -> bool { diagnostics .iter() diff --git a/crates/core/tests/unit/dynamic_worker_tests.rs b/crates/core/tests/unit/dynamic_worker_tests.rs index cbe7c4519..32c5dccf5 100644 --- a/crates/core/tests/unit/dynamic_worker_tests.rs +++ b/crates/core/tests/unit/dynamic_worker_tests.rs @@ -476,6 +476,7 @@ fn registration_plan_and_scope_type_helpers_validate_edges() { surface: RegistrationSurface::Subscriber as i32, priority: 0, break_chain: false, + export_target: None, }], error: None, }, @@ -491,6 +492,7 @@ fn registration_plan_and_scope_type_helpers_validate_edges() { surface: 999, priority: 0, break_chain: false, + export_target: None, }], error: None, }, @@ -510,6 +512,7 @@ fn registration_plan_and_scope_type_helpers_validate_edges() { surface: RegistrationSurface::Unspecified as i32, priority: 0, break_chain: false, + export_target: None, }], error: None, }, @@ -561,6 +564,61 @@ fn registration_plan_and_scope_type_helpers_validate_edges() { assert_eq!(proto_scope_type(999), crate::api::scope::ScopeType::Custom); } +#[test] +fn unadvertised_export_activation_hook_registration_is_ignored() { + let policy = registration( + RegistrationSurface::ExportActivationPolicy, + "export_activation_policy", + ); + let subscriber = registration(RegistrationSurface::Subscriber, "subscriber"); + let mut registrations = vec![policy.clone(), subscriber.clone()]; + + filter_unadvertised_export_activation_hooks( + "fixture_worker", + true, + &[RegistrationSurface::Subscriber as i32], + &mut registrations, + ); + assert_eq!(registrations, vec![subscriber]); + + let mut advertised = vec![policy.clone()]; + filter_unadvertised_export_activation_hooks( + "fixture_worker", + true, + &[RegistrationSurface::ExportActivationPolicy as i32], + &mut advertised, + ); + assert_eq!(advertised, vec![policy.clone()]); + + let target = registration(RegistrationSurface::ExportTarget, "exporter"); + let mut unadvertised_target = vec![target.clone()]; + filter_unadvertised_export_activation_hooks( + "fixture_worker", + true, + &[RegistrationSurface::ExportActivationPolicy as i32], + &mut unadvertised_target, + ); + assert!(unadvertised_target.is_empty()); + + let mut advertised_target = vec![target.clone()]; + filter_unadvertised_export_activation_hooks( + "fixture_worker", + true, + &[RegistrationSurface::ExportTarget as i32], + &mut advertised_target, + ); + assert_eq!(advertised_target, vec![target]); + + let mut undeclared_policy = vec![policy]; + filter_unadvertised_export_activation_hooks( + "fixture_worker", + false, + &[RegistrationSurface::ExportActivationPolicy as i32], + &mut undeclared_policy, + ); + assert!(undeclared_policy.is_empty()); +} + #[test] fn relay_compatibility_and_blocking_helpers_cover_local_edges() { enable_operational_logs(); @@ -2892,6 +2950,7 @@ fn registration(surface: RegistrationSurface, local_name: &str) -> Registration surface: surface as i32, priority: 0, break_chain: false, + export_target: None, } } diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 711f60f15..fa62c61ba 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -504,6 +504,7 @@ fn native_test_adapter( plugin_kind: "test.native.adapter".into(), relay_compat: "^0.8".into(), allows_multiple_components: false, + allows_export_activation_policy: true, plugin: Mutex::new(plugin), _library: libloading::os::unix::Library::this().into(), }), @@ -4694,6 +4695,7 @@ fn native_registration_entrypoints_reject_invalid_host_contexts_and_names() { plugin_kind: "test.native".into(), relay_compat: "^0.8".into(), allows_multiple_components: false, + allows_export_activation_policy: true, plugin: Mutex::new(NemoRelayNativePluginV1::default()), _library: libloading::os::unix::Library::this().into(), }); @@ -5130,6 +5132,7 @@ fn assert_async_request_registration_rejects_legacy_relay_contract() { plugin_kind: "test.native.legacy".into(), relay_compat: "^0.5".into(), allows_multiple_components: false, + allows_export_activation_policy: false, plugin: Mutex::new(NemoRelayNativePluginV1::default()), _library: libloading::os::unix::Library::this().into(), }); @@ -5179,6 +5182,7 @@ async fn native_async_wrappers_validate_callback_result_shapes() { plugin_kind: "test.native.async".into(), relay_compat: "^0.8".into(), allows_multiple_components: false, + allows_export_activation_policy: true, plugin: Mutex::new(NemoRelayNativePluginV1::default()), _library: libloading::os::unix::Library::this().into(), }); @@ -6044,6 +6048,7 @@ async fn native_callback_wrappers_release_error_outputs_and_preserve_reasons() { plugin_kind: "test.native.callback-errors".into(), relay_compat: "^0.8".into(), allows_multiple_components: false, + allows_export_activation_policy: true, plugin: Mutex::new(NemoRelayNativePluginV1::default()), _library: libloading::os::unix::Library::this().into(), }); diff --git a/crates/core/tests/unit/observability/plugin_component_tests.rs b/crates/core/tests/unit/observability/plugin_component_tests.rs index 0446c8ee7..80f559060 100644 --- a/crates/core/tests/unit/observability/plugin_component_tests.rs +++ b/crates/core/tests/unit/observability/plugin_component_tests.rs @@ -17,7 +17,7 @@ use crate::config_editor::{EditorConfig, EditorFieldKind, EditorSchema}; use crate::plugin::plugin_config_schema; use crate::plugin::{ PluginComponentSpec, PluginConfig, clear_plugin_configuration, initialize_plugins_exact, - list_plugin_kinds, lookup_plugin, validate_plugin_config, + list_plugin_kinds, lookup_plugin, rollback_registrations, validate_plugin_config, }; use serde_json::json; use std::fs; @@ -573,6 +573,7 @@ fn signal_endpoint_resolution_derives_or_preserves_the_expected_destination() { service_version: None, instrumentation_scope: default_otel_instrumentation_scope(), timeout_millis: default_timeout_millis(), + activation_policy: None, }]; assert_eq!( resolve_signal_endpoints("metrics", Some(&custom), &[]).unwrap()[0].endpoint, @@ -598,6 +599,7 @@ fn signal_endpoint_resolution_rejects_ambiguous_trace_paths_and_wrong_signal_pat service_version: None, instrumentation_scope: default_otel_instrumentation_scope(), timeout_millis: default_timeout_millis(), + activation_policy: None, }; assert!(resolve_signal_endpoints("logs", Some(&vec![explicit.clone()]), &[]).is_err()); explicit.endpoint = "https://collector.example/".to_string(); @@ -612,9 +614,18 @@ fn signal_endpoint_resolution_covers_missing_trace_grpc_and_invalid_transports() let mut grpc_trace = test_opentelemetry_endpoint(); grpc_trace.transport = "grpc".to_string(); grpc_trace.endpoint = "http://collector.example:4317".to_string(); + grpc_trace.activation_policy = Some(ExportActivationPolicyConfig { + provider: "test.inherited-policy".into(), + timeout_millis: 1_234, + config: json!({"consent": true}), + }); let derived = resolve_signal_endpoints("metrics", None, &[grpc_trace]).unwrap(); assert_eq!(derived[0].endpoint, "http://collector.example:4317"); assert_eq!(derived[0].transport, "grpc"); + assert_eq!( + derived[0].activation_policy.as_ref().unwrap().provider, + "test.inherited-policy" + ); for endpoint in ["ftp://collector.example", "not a url"] { let mut trace = test_opentelemetry_endpoint(); @@ -776,6 +787,7 @@ fn default_config_and_component_conversion_cover_public_shape() { mark_exclude_names: default_mark_exclude_names(), attribute_mappings: Vec::new(), promote_metadata_prefixes: Vec::new(), + activation_policy: None, }], logs: None, metrics: None, @@ -963,6 +975,7 @@ fn opentelemetry_endpoint_header_env_is_resolved_and_snapshotted() { mark_exclude_names: default_mark_exclude_names(), attribute_mappings: Vec::new(), promote_metadata_prefixes: Vec::new(), + activation_policy: None, }, ) .unwrap(); @@ -991,6 +1004,7 @@ fn test_opentelemetry_endpoint() -> OpenTelemetryEndpointConfig { mark_exclude_names: default_mark_exclude_names(), attribute_mappings: Vec::new(), promote_metadata_prefixes: Vec::new(), + activation_policy: None, } } @@ -1006,6 +1020,7 @@ fn test_signal_endpoint() -> OpenTelemetrySignalEndpointConfig { service_version: None, instrumentation_scope: default_otel_instrumentation_scope(), timeout_millis: default_timeout_millis(), + activation_policy: None, } } @@ -1182,8 +1197,8 @@ fn validate_opentelemetry_section_reports_empty_and_malformed_endpoints() { } } -#[test] -fn opentelemetry_registration_rejects_an_empty_endpoint_list() { +#[tokio::test] +async fn opentelemetry_registration_rejects_an_empty_endpoint_list() { let mut context = PluginRegistrationContext::new(); let error = register_opentelemetry( OpenTelemetrySectionConfig { @@ -1194,6 +1209,7 @@ fn opentelemetry_registration_rejects_an_empty_endpoint_list() { }, &mut context, ) + .await .unwrap_err(); assert!(error.to_string().contains("at least one endpoint")); } @@ -2044,6 +2060,7 @@ fn build_atof_sink_config_maps_headers_timeout_and_rejects_transport() { )]), timeout_millis: 123, field_name_policy: "replace_dots".into(), + activation_policy: None, }), ) .unwrap(); @@ -2077,6 +2094,7 @@ fn build_atof_sink_config_maps_headers_timeout_and_rejects_transport() { header_env: std::collections::HashMap::new(), timeout_millis: 3_000, field_name_policy: "preserve".into(), + activation_policy: None, }), ) .unwrap_err(); @@ -2092,6 +2110,7 @@ fn build_atof_sink_config_maps_headers_timeout_and_rejects_transport() { header_env: std::collections::HashMap::new(), timeout_millis: 3_000, field_name_policy: "bogus".into(), + activation_policy: None, }), ) .unwrap_err(); @@ -2315,6 +2334,7 @@ fn atif_remote_storage_validates_s3_configuration_and_http_access_outcomes() { region: Some("us-east-1".into()), endpoint_url: Some("http://127.0.0.1:9".into()), allow_http: Some(true), + activation_policy: None, }); let storage = build_atif_storage(3, &s3).expect("S3 client configuration should resolve"); drop(storage); @@ -2330,6 +2350,7 @@ fn atif_remote_storage_validates_s3_configuration_and_http_access_outcomes() { headers: std::collections::HashMap::new(), header_env: std::collections::HashMap::new(), timeout_millis: 5_000, + activation_policy: None, }), ) .unwrap(); @@ -4792,6 +4813,7 @@ fn atif_storage_private_helpers_resolve_env_and_key_prefix_branches() { region: Some("us-west-2".into()), endpoint_url: Some("http://127.0.0.1:9000".into()), allow_http: Some(true), + activation_policy: None, }, ) .unwrap(); @@ -4821,6 +4843,7 @@ fn http_storage_config(endpoint: impl Into) -> HttpStorageConfig { headers: std::collections::HashMap::new(), header_env: std::collections::HashMap::new(), timeout_millis: 1_000, + activation_policy: None, } } @@ -4906,6 +4929,7 @@ fn s3_remote_storage_uploads_to_a_custom_http_endpoint() { region: Some("us-east-1".into()), endpoint_url: Some(endpoint), allow_http: Some(true), + activation_policy: None, }), ) .unwrap(); @@ -4977,3 +5001,438 @@ fn atif_filename_helpers_cover_metadata_resolution_and_rejection_paths() { .is_err() ); } + +#[tokio::test] +#[allow(clippy::await_holding_lock)] +async fn export_activation_policy_is_evaluated_once_before_otel_construction() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + let calls = Arc::new(AtomicUsize::new(0)); + let callback_calls = Arc::clone(&calls); + let mut provider_context = PluginRegistrationContext::new(); + provider_context + .register_export_activation_policy_for_provider( + "test.allow-export", + Arc::new(move |request| { + let callback_calls = Arc::clone(&callback_calls); + Box::pin(async move { + callback_calls.fetch_add(1, Ordering::SeqCst); + assert_eq!(request.target_kind, ExportActivationTargetKind::OtlpTrace); + assert_eq!(request.config, json!({"region": "local"})); + Ok(ExportActivationDecision::Allow) + }) + }), + ) + .unwrap(); + + let mut endpoint = test_opentelemetry_endpoint(); + endpoint.activation_policy = Some(ExportActivationPolicyConfig { + provider: "test.allow-export".into(), + timeout_millis: 5_000, + config: json!({"region": "local"}), + }); + let mut observability_context = PluginRegistrationContext::with_export_activation_policies( + provider_context.export_activation_policies(), + ); + register_opentelemetry( + OpenTelemetrySectionConfig { + enabled: true, + endpoints: vec![endpoint], + logs: None, + metrics: None, + }, + &mut observability_context, + ) + .await + .unwrap(); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + let mut observability_registrations = observability_context.into_registrations(); + rollback_registrations(&mut observability_registrations); + let mut provider_registrations = provider_context.into_registrations(); + rollback_registrations(&mut provider_registrations); +} + +#[tokio::test] +#[allow(clippy::await_holding_lock)] +async fn unavailable_and_timed_out_export_policies_fail_closed() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + let mut unavailable = test_opentelemetry_endpoint(); + unavailable.activation_policy = Some(ExportActivationPolicyConfig { + provider: "test.missing-export-policy".into(), + timeout_millis: 5_000, + config: Json::Null, + }); + let mut context = PluginRegistrationContext::new(); + register_opentelemetry( + OpenTelemetrySectionConfig { + enabled: true, + endpoints: vec![unavailable], + logs: None, + metrics: None, + }, + &mut context, + ) + .await + .unwrap(); + assert!(context.into_registrations().is_empty()); + + let mut provider_context = PluginRegistrationContext::new(); + provider_context + .register_export_activation_policy_for_provider( + "test.slow-export-policy", + Arc::new(|_| { + Box::pin(async { + tokio::time::sleep(Duration::from_millis(1_100)).await; + Ok(ExportActivationDecision::Allow) + }) + }), + ) + .unwrap(); + let mut timed_out = test_opentelemetry_endpoint(); + timed_out.activation_policy = Some(ExportActivationPolicyConfig { + provider: "test.slow-export-policy".into(), + timeout_millis: 1_000, + config: Json::Null, + }); + let mut context = PluginRegistrationContext::with_export_activation_policies( + provider_context.export_activation_policies(), + ); + register_opentelemetry( + OpenTelemetrySectionConfig { + enabled: true, + endpoints: vec![timed_out], + logs: None, + metrics: None, + }, + &mut context, + ) + .await + .unwrap(); + assert!(context.into_registrations().is_empty()); + let mut provider_registrations = provider_context.into_registrations(); + rollback_registrations(&mut provider_registrations); +} + +#[tokio::test] +#[allow(clippy::await_holding_lock)] +async fn denied_and_failed_export_policies_fail_closed() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + for (provider, result) in [ + ( + "test.deny-export-policy", + Ok(ExportActivationDecision::Deny), + ), + ( + "test.failed-export-policy", + Err(FlowError::Internal("policy failed".into())), + ), + ] { + let mut provider_context = PluginRegistrationContext::new(); + let callback_result = result.clone(); + provider_context + .register_export_activation_policy_for_provider( + provider, + Arc::new(move |_| { + let callback_result = callback_result.clone(); + Box::pin(async move { callback_result }) + }), + ) + .unwrap(); + + let mut endpoint = test_opentelemetry_endpoint(); + endpoint.activation_policy = Some(ExportActivationPolicyConfig { + provider: provider.into(), + timeout_millis: 5_000, + config: Json::Null, + }); + let mut context = PluginRegistrationContext::with_export_activation_policies( + provider_context.export_activation_policies(), + ); + register_opentelemetry( + OpenTelemetrySectionConfig { + enabled: true, + endpoints: vec![endpoint], + logs: None, + metrics: None, + }, + &mut context, + ) + .await + .unwrap(); + assert!(context.into_registrations().is_empty()); + let mut provider_registrations = provider_context.into_registrations(); + rollback_registrations(&mut provider_registrations); + } +} + +#[test] +fn export_activation_provider_uniqueness_is_scoped_to_one_activation() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + let callback = Arc::new(|_| Box::pin(async { Ok(ExportActivationDecision::Allow) }) as _); + let mut first = PluginRegistrationContext::new(); + first + .register_export_activation_policy_for_provider("test.duplicate-policy", callback.clone()) + .unwrap(); + let registry = first.export_activation_policies(); + let mut duplicate = PluginRegistrationContext::with_export_activation_policies(registry); + assert!( + duplicate + .register_export_activation_policy_for_provider( + "test.duplicate-policy", + callback.clone(), + ) + .is_err() + ); + let mut independent_activation = PluginRegistrationContext::new(); + independent_activation + .register_export_activation_policy_for_provider("test.duplicate-policy", callback.clone()) + .unwrap(); + + let mut first_registrations = first.into_registrations(); + rollback_registrations(&mut first_registrations); + duplicate + .register_export_activation_policy_for_provider("test.duplicate-policy", callback) + .unwrap(); + let mut duplicate_registrations = duplicate.into_registrations(); + rollback_registrations(&mut duplicate_registrations); + let mut independent_registrations = independent_activation.into_registrations(); + rollback_registrations(&mut independent_registrations); +} + +#[tokio::test] +#[allow(clippy::await_holding_lock)] +async fn all_denied_atif_storage_does_not_fall_back_to_local_files() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + let section = AtifSectionConfig { + enabled: true, + output_directory: Some(temp_dir("denied-atif-fallback")), + storage: vec![AtifStorageConfig::Http(HttpStorageConfig { + endpoint: "https://collector.example/atif".into(), + headers: HashMap::new(), + header_env: HashMap::new(), + timeout_millis: 3_000, + activation_policy: Some(ExportActivationPolicyConfig { + provider: "test.missing-atif-policy".into(), + timeout_millis: 5_000, + config: Json::Null, + }), + })], + ..AtifSectionConfig::default() + }; + let mut context = PluginRegistrationContext::new(); + register_atif_dispatcher(section, &mut context) + .await + .unwrap(); + assert!(context.into_registrations().is_empty()); +} + +#[tokio::test] +#[allow(clippy::await_holding_lock)] +async fn denied_atof_stream_retains_allowed_local_sink() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + reset_runtime(); + let output_directory = temp_dir("allowed-local-atof"); + let section = AtofSectionConfig { + enabled: true, + sinks: vec![ + AtofSinkSectionConfig::File(AtofFileSinkSectionConfig { + output_directory: Some(output_directory.clone()), + filename: Some("events.jsonl".into()), + mode: "overwrite".into(), + activation_policy: None, + }), + AtofSinkSectionConfig::Stream(AtofStreamSinkSectionConfig { + name: Some("denied".into()), + url: "https://collector.example/events".into(), + // Construction would fail if policy filtering retained this sink. + transport: "invalid".into(), + headers: HashMap::new(), + header_env: HashMap::new(), + timeout_millis: 3_000, + field_name_policy: "preserve".into(), + activation_policy: Some(ExportActivationPolicyConfig { + provider: "test.missing-atof-policy".into(), + timeout_millis: 5_000, + config: Json::Null, + }), + }), + ], + }; + let mut context = PluginRegistrationContext::new(); + register_atof_exporter(section, &mut context).await.unwrap(); + let mut registrations = context.into_registrations(); + let agent = push_agent("allowed-local-atof-agent"); + crate::api::scope::event( + crate::api::scope::EmitMarkEventParams::builder() + .name("checkpoint") + .parent(&agent) + .build(), + ) + .unwrap(); + pop(&agent); + rollback_registrations(&mut registrations); + + let content = fs::read_to_string(output_directory.join("events.jsonl")).unwrap(); + assert_eq!(content.lines().count(), 3); + assert!(content.contains("\"name\":\"checkpoint\"")); +} + +#[tokio::test] +#[allow(clippy::await_holding_lock)] +async fn denied_atof_file_sink_creates_no_file_or_subscriber() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + let output_directory = temp_dir("denied-local-atof"); + let section = AtofSectionConfig { + enabled: true, + sinks: vec![AtofSinkSectionConfig::File(AtofFileSinkSectionConfig { + output_directory: Some(output_directory.clone()), + filename: Some("events.jsonl".into()), + mode: "overwrite".into(), + activation_policy: Some(ExportActivationPolicyConfig { + provider: "test.missing-local-atof-policy".into(), + timeout_millis: 5_000, + config: Json::Null, + }), + })], + }; + let mut context = PluginRegistrationContext::new(); + register_atof_exporter(section, &mut context).await.unwrap(); + assert!(context.into_registrations().is_empty()); + assert!(!output_directory.join("events.jsonl").exists()); +} + +#[tokio::test] +#[allow(clippy::await_holding_lock)] +async fn denied_atif_local_destination_disables_local_only_dispatcher() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + let output_directory = temp_dir("denied-local-atif"); + let section = AtifSectionConfig { + enabled: true, + output_directory: Some(output_directory.clone()), + local_activation_policy: Some(ExportActivationPolicyConfig { + provider: "test.missing-local-atif-policy".into(), + timeout_millis: 5_000, + config: Json::Null, + }), + ..AtifSectionConfig::default() + }; + let mut context = PluginRegistrationContext::new(); + register_atif_dispatcher(section, &mut context) + .await + .unwrap(); + assert!(context.into_registrations().is_empty()); + assert_eq!(fs::read_dir(output_directory).unwrap().count(), 0); +} + +#[tokio::test] +#[allow(clippy::await_holding_lock)] +async fn denied_atof_stream_preserves_later_sink_index() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + let section = AtofSectionConfig { + enabled: true, + sinks: vec![ + AtofSinkSectionConfig::Stream(AtofStreamSinkSectionConfig { + name: Some("denied".into()), + url: "https://collector.example/denied".into(), + transport: "http_post".into(), + headers: HashMap::new(), + header_env: HashMap::new(), + timeout_millis: 3_000, + field_name_policy: "preserve".into(), + activation_policy: Some(ExportActivationPolicyConfig { + provider: "test.missing-atof-policy".into(), + timeout_millis: 30_000, + config: Json::Null, + }), + }), + AtofSinkSectionConfig::Stream(AtofStreamSinkSectionConfig { + name: Some("invalid".into()), + url: "https://collector.example/invalid".into(), + transport: "invalid".into(), + headers: HashMap::new(), + header_env: HashMap::new(), + timeout_millis: 3_000, + field_name_policy: "preserve".into(), + activation_policy: None, + }), + ], + }; + let error = register_atof_exporter(section, &mut PluginRegistrationContext::new()) + .await + .unwrap_err(); + assert!(error.to_string().contains("sinks[1]"), "{error}"); +} + +#[test] +fn filtered_atif_storage_preserves_original_sink_labels() { + let dispatcher = AtifDispatcher::with_remote_storage_indices( + AtifSectionConfig { + storage: vec![AtifStorageConfig::Http(HttpStorageConfig { + endpoint: "https://collector.example/atif".into(), + headers: HashMap::new(), + header_env: HashMap::new(), + timeout_millis: 3_000, + activation_policy: None, + })], + ..AtifSectionConfig::default() + }, + vec![2], + true, + ); + assert_eq!(dispatcher.sink_targets(), vec![SinkLabel::Remote(2)]); +} + +#[test] +fn export_activation_policy_config_validates_provider_and_timeout() { + let config = json!({ + "version": 4, + "opentelemetry": { + "enabled": true, + "endpoints": [ + { + "type": "full", + "endpoint": "http://localhost:4318/v1/traces", + "activation_policy": { + "provider": " ", + "timeout_millis": 300001 + } + }, + { + "type": "full", + "endpoint": "http://localhost:4318/v1/traces", + "activation_policy": { + "provider": "test.policy", + "timeout_millis": 999 + } + } + ] + } + }); + let diagnostics = validate_observability_plugin_config(config.as_object().unwrap()); + assert!(diagnostics.iter().any(|diagnostic| { + diagnostic.field.as_deref() == Some("opentelemetry.endpoints[0].activation_policy.provider") + })); + assert!(diagnostics.iter().any(|diagnostic| { + diagnostic.field.as_deref() + == Some("opentelemetry.endpoints[0].activation_policy.timeout_millis") + })); + assert!(diagnostics.iter().any(|diagnostic| { + diagnostic.field.as_deref() + == Some("opentelemetry.endpoints[1].activation_policy.timeout_millis") + })); +} + +#[test] +fn export_activation_policy_timeout_is_clamped_at_evaluation() { + let defaulted: ExportActivationPolicyConfig = serde_json::from_value(json!({ + "provider": "test.policy" + })) + .unwrap(); + assert_eq!(defaulted.timeout_millis, 30_000); + assert_eq!(export_activation_timeout(0), Duration::from_secs(1)); + assert_eq!(export_activation_timeout(30_000), Duration::from_secs(30)); + assert_eq!( + export_activation_timeout(u64::MAX), + Duration::from_secs(300) + ); +} diff --git a/crates/core/tests/unit/plugin_dynamic_host_tests.rs b/crates/core/tests/unit/plugin_dynamic_host_tests.rs index f959b8eb2..c7343439a 100644 --- a/crates/core/tests/unit/plugin_dynamic_host_tests.rs +++ b/crates/core/tests/unit/plugin_dynamic_host_tests.rs @@ -26,6 +26,7 @@ fn unsafe_kind_deregistration_retains_runtime_and_owner() { let owner_id = claim.owner_id(); let mut activation = PluginHostActivation { active: true, + export_activation_policies: None, native: Some(NativePluginActivation::with_plugin_kind_for_test( "fixture.poisoned", )), diff --git a/crates/core/tests/unit/plugin_dynamic_tests.rs b/crates/core/tests/unit/plugin_dynamic_tests.rs index 071182d47..23de992e7 100644 --- a/crates/core/tests/unit/plugin_dynamic_tests.rs +++ b/crates/core/tests/unit/plugin_dynamic_tests.rs @@ -628,6 +628,7 @@ enabled = false items = [ "plugin_worker", "config_schema", + "export_activation_policy", ] [config_schema] @@ -645,6 +646,7 @@ entrypoint = "acme_guardrails.plugin:register" vec![ DynamicPluginCapability::PluginWorker, DynamicPluginCapability::ConfigSchema, + DynamicPluginCapability::ExportActivationPolicy, ] ); assert_eq!( diff --git a/crates/core/tests/unit/plugin_tests.rs b/crates/core/tests/unit/plugin_tests.rs index 7816bfbb5..cffe24a88 100644 --- a/crates/core/tests/unit/plugin_tests.rs +++ b/crates/core/tests/unit/plugin_tests.rs @@ -1735,8 +1735,12 @@ fn test_pending_rollbacks_ignore_delivery_only_errors() { )]); } { - let mut pending = - PendingPluginRegistrationContext::new("fixture.".into(), Some(Arc::clone(&failures))); + let mut pending = PendingPluginRegistrationContext::new( + "fixture.".into(), + "fixture".into(), + Some(Arc::clone(&failures)), + Arc::new(ExportActivationPolicyRegistry::default()), + ); pending .context .add_registration(PluginRegistration::new_with_outcome( diff --git a/crates/ffi/nemo_relay.h b/crates/ffi/nemo_relay.h index 0d2bc97ef..8d8e14be8 100644 --- a/crates/ffi/nemo_relay.h +++ b/crates/ffi/nemo_relay.h @@ -406,6 +406,16 @@ typedef NemoRelayStatus (*NemoRelayPluginRegisterCb)(void *user_data, const char *plugin_config_json, struct FfiPluginContext *ctx); +/** + * Callback for an activation-scoped export policy provider. + */ +typedef char *(*NemoRelayExportActivationPolicyCb)(void *user_data, const char *request_json); + +/** + * One-shot callback that constructs an allowed export target. + */ +typedef NemoRelayStatus (*NemoRelayExportTargetActivationCb)(void *user_data); + /** * Callback for tool request/response sanitization guardrails and intercepts. * Receives tool name and arguments as JSON, returns sanitized arguments as JSON. @@ -2057,6 +2067,30 @@ NemoRelayStatus nemo_relay_plugin_context_register_subscriber(struct FfiPluginCo void *user_data, NemoRelayFreeFn free_fn); +/** + * Register the activation-scoped export policy owned by this plugin component. + * + * # Safety + * `ctx` must be valid and callback user data must remain valid until `free_fn` runs. + */ +NemoRelayStatus nemo_relay_plugin_context_register_export_activation_policy(struct FfiPluginContext *ctx, + NemoRelayExportActivationPolicyCb cb, + void *user_data, + NemoRelayFreeFn free_fn); + +/** + * Register one deferred local or remote export target. + * + * # Safety + * `ctx` and `registration_json` must be valid and callback user data must + * remain valid until `free_fn` runs. + */ +NemoRelayStatus nemo_relay_plugin_context_register_export_target(struct FfiPluginContext *ctx, + const char *registration_json, + NemoRelayExportTargetActivationCb cb, + void *user_data, + NemoRelayFreeFn free_fn); + /** * Register a mark event sanitizer into a plugin context. * # Safety diff --git a/crates/ffi/src/api/mod.rs b/crates/ffi/src/api/mod.rs index d201a7475..33e9590f3 100644 --- a/crates/ffi/src/api/mod.rs +++ b/crates/ffi/src/api/mod.rs @@ -15,16 +15,19 @@ use std::time::Duration; use crate::callable::{ NemoRelayCodecDecodeFn, NemoRelayCodecEncodeFn, NemoRelayCollectorCb, NemoRelayEventSanitizeCb, - NemoRelayEventSubscriberCb, NemoRelayFinalizerCb, NemoRelayFreeFn, NemoRelayLlmConditionalCb, - NemoRelayLlmExecCb, NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestInterceptCb, - NemoRelayLlmSanitizeRequestCb, NemoRelayLlmSanitizeResponseCb, NemoRelayPluginRegisterCb, - NemoRelayPluginValidateCb, NemoRelayToolConditionalCb, NemoRelayToolExecCb, - NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, wrap_codec_fn, wrap_collector_fn, - wrap_event_sanitize_fn, wrap_event_subscriber, wrap_finalizer_fn, wrap_llm_conditional_fn, - wrap_llm_exec_fn, wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, - wrap_llm_sanitize_request_fn, wrap_llm_sanitize_response_fn, wrap_llm_stream_exec_fn, - wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, wrap_tool_exec_fn, - wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, + NemoRelayEventSubscriberCb, NemoRelayExportActivationPolicyCb, + NemoRelayExportTargetActivationCb, NemoRelayFinalizerCb, NemoRelayFreeFn, + NemoRelayLlmConditionalCb, NemoRelayLlmExecCb, NemoRelayLlmExecInterceptCb, + NemoRelayLlmRequestInterceptCb, NemoRelayLlmSanitizeRequestCb, NemoRelayLlmSanitizeResponseCb, + NemoRelayPluginRegisterCb, NemoRelayPluginValidateCb, NemoRelayToolConditionalCb, + NemoRelayToolExecCb, NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, wrap_codec_fn, + wrap_collector_fn, wrap_event_sanitize_fn, wrap_event_subscriber, + wrap_export_activation_policy_fn, wrap_export_target_activation_fn, wrap_finalizer_fn, + wrap_llm_conditional_fn, wrap_llm_exec_fn, wrap_llm_exec_intercept_fn, + wrap_llm_request_intercept_fn, wrap_llm_sanitize_request_fn, wrap_llm_sanitize_response_fn, + wrap_llm_stream_exec_fn, wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, + wrap_tool_exec_fn, wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, + wrap_tool_sanitize_fn, }; use crate::convert::{ c_str_to_json, c_str_to_opt_json, c_str_to_string, json_to_c_string, nemo_relay_string_free, diff --git a/crates/ffi/src/api/plugin.rs b/crates/ffi/src/api/plugin.rs index 09cd151f8..fcec59098 100644 --- a/crates/ffi/src/api/plugin.rs +++ b/crates/ffi/src/api/plugin.rs @@ -4,21 +4,24 @@ use super::{ Arc, CStr, ConfigDiagnostic, DiagnosticLevel, DynamicPluginActivationSpec, FfiPluginActivation, FfiPluginContext, Future, NemoRelayEventSanitizeCb, NemoRelayEventSubscriberCb, - NemoRelayFreeFn, NemoRelayLlmConditionalCb, NemoRelayLlmExecInterceptCb, - NemoRelayLlmRequestInterceptCb, NemoRelayLlmSanitizeRequestCb, NemoRelayLlmSanitizeResponseCb, - NemoRelayPluginRegisterCb, NemoRelayPluginValidateCb, NemoRelayStatus, - NemoRelayToolConditionalCb, NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, Pin, Plugin, - PluginConfig, PluginError, PluginHostActivation, PluginRegistrationContext, - active_plugin_report, c_char, c_str_to_json, c_str_to_string, clear_last_error, - clear_plugin_configuration, deregister_plugin, initialize_plugins, json_to_c_string, - last_error_message, list_plugin_kinds, nemo_relay_string_free, register_adaptive_component, - register_plugin, set_last_error, status_from_plugin_error, tokio_runtime, - validate_plugin_config, wrap_event_sanitize_fn, wrap_event_subscriber, wrap_llm_conditional_fn, - wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, wrap_llm_sanitize_request_fn, - wrap_llm_sanitize_response_fn, wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, - wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, + NemoRelayExportActivationPolicyCb, NemoRelayExportTargetActivationCb, NemoRelayFreeFn, + NemoRelayLlmConditionalCb, NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestInterceptCb, + NemoRelayLlmSanitizeRequestCb, NemoRelayLlmSanitizeResponseCb, NemoRelayPluginRegisterCb, + NemoRelayPluginValidateCb, NemoRelayStatus, NemoRelayToolConditionalCb, + NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, Pin, Plugin, PluginConfig, PluginError, + PluginHostActivation, PluginRegistrationContext, active_plugin_report, c_char, c_str_to_json, + c_str_to_string, clear_last_error, clear_plugin_configuration, deregister_plugin, + initialize_plugins, json_to_c_string, last_error_message, list_plugin_kinds, + nemo_relay_string_free, register_adaptive_component, register_plugin, set_last_error, + status_from_plugin_error, tokio_runtime, validate_plugin_config, wrap_event_sanitize_fn, + wrap_event_subscriber, wrap_export_activation_policy_fn, wrap_export_target_activation_fn, + wrap_llm_conditional_fn, wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, + wrap_llm_sanitize_request_fn, wrap_llm_sanitize_response_fn, wrap_llm_stream_exec_intercept_fn, + wrap_tool_conditional_fn, wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, + wrap_tool_sanitize_fn, }; use crate::api::event_registry::Surface; +use nemo_relay::api::export_activation::ExportTargetRegistration; use nemo_relay_pii_redaction::component::register_pii_redaction_component; struct FfiHostedPluginUserData { @@ -542,6 +545,65 @@ pub unsafe extern "C" fn nemo_relay_plugin_context_register_subscriber( } } +/// Register the activation-scoped export policy owned by this plugin component. +/// +/// # Safety +/// `ctx` must be valid and callback user data must remain valid until `free_fn` runs. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_plugin_context_register_export_activation_policy( + ctx: *mut FfiPluginContext, + cb: NemoRelayExportActivationPolicyCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> NemoRelayStatus { + clear_last_error(); + if ctx.is_null() { + set_last_error("plugin context is null"); + return NemoRelayStatus::NullPointer; + } + let callback = wrap_export_activation_policy_fn(cb, user_data, free_fn); + match unsafe { &mut *((*ctx).0) }.register_export_activation_policy(callback) { + Ok(()) => NemoRelayStatus::Ok, + Err(error) => status_from_plugin_error(&error), + } +} + +/// Register one deferred local or remote export target. +/// +/// # Safety +/// `ctx` and `registration_json` must be valid and callback user data must +/// remain valid until `free_fn` runs. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_plugin_context_register_export_target( + ctx: *mut FfiPluginContext, + registration_json: *const c_char, + cb: NemoRelayExportTargetActivationCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> NemoRelayStatus { + clear_last_error(); + if ctx.is_null() { + set_last_error("plugin context is null"); + return NemoRelayStatus::NullPointer; + } + let Some(value) = c_str_to_json(registration_json) else { + set_last_error("export target registration JSON is null or invalid"); + return NemoRelayStatus::InvalidArg; + }; + let registration = match serde_json::from_value::(value) { + Ok(registration) => registration, + Err(error) => { + set_last_error(&format!("invalid export target registration: {error}")); + return NemoRelayStatus::InvalidArg; + } + }; + let activate = wrap_export_target_activation_fn(cb, user_data, free_fn); + match unsafe { &mut *((*ctx).0) }.register_export_target(registration, activate) { + Ok(()) => NemoRelayStatus::Ok, + Err(error) => status_from_plugin_error(&error), + } +} + unsafe fn plugin_register_event_sanitizer( ctx: *mut FfiPluginContext, name: *const c_char, diff --git a/crates/ffi/src/callable.rs b/crates/ffi/src/callable.rs index 280b9dfb1..e926839c1 100644 --- a/crates/ffi/src/callable.rs +++ b/crates/ffi/src/callable.rs @@ -32,6 +32,9 @@ use serde_json::Value as Json; use tokio_stream::StreamExt; use nemo_relay::api::event::{Event, EventSanitizeFields}; +use nemo_relay::api::export_activation::{ + ExportActivationDecision, ExportActivationPolicyFn, ExportActivationRequest, +}; use nemo_relay::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; use nemo_relay::api::tool::{ToolExecutionInterceptOutcome, ToolExecutionResult}; use nemo_relay::codec::request::AnnotatedLlmRequest as AnnotatedLLMRequest; @@ -292,6 +295,14 @@ pub type NemoRelayPluginRegisterCb = unsafe extern "C" fn( ctx: *mut FfiPluginContext, ) -> NemoRelayStatus; +/// Callback for an activation-scoped export policy provider. +pub type NemoRelayExportActivationPolicyCb = + unsafe extern "C" fn(user_data: *mut libc::c_void, request_json: *const c_char) -> *mut c_char; + +/// One-shot callback that constructs an allowed export target. +pub type NemoRelayExportTargetActivationCb = + unsafe extern "C" fn(user_data: *mut libc::c_void) -> NemoRelayStatus; + // --------------------------------------------------------------------------- // Shared user_data wrapper (ensures cleanup) // --------------------------------------------------------------------------- @@ -324,6 +335,62 @@ fn make_user_data( }) } +/// Wraps an export activation policy callback. +pub fn wrap_export_activation_policy_fn( + cb: NemoRelayExportActivationPolicyCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> ExportActivationPolicyFn { + let ud = make_user_data(user_data, free_fn); + Arc::new(move |request: ExportActivationRequest| { + let ud = Arc::clone(&ud); + Box::pin(async move { + clear_last_error(); + let request = json_to_c_string(&serde_json::to_value(request).map_err(|error| { + FlowError::Internal(format!( + "failed to serialize export activation request: {error}" + )) + })?); + let result_ptr = unsafe { cb(ud.ptr, request) }; + unsafe { nemo_relay_string_free_internal(request) }; + let result = json_result_from_ptr( + result_ptr, + "export activation policy callback returned null", + ) + .and_then(|value| { + serde_json::from_value::(value).map_err(|error| { + FlowError::Internal(format!( + "invalid export activation policy decision: {error}" + )) + }) + }); + unsafe { nemo_relay_string_free_internal(result_ptr) }; + result + }) + }) +} + +/// Wraps a one-shot export-target activation callback. +pub fn wrap_export_target_activation_fn( + cb: NemoRelayExportTargetActivationCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> impl Fn() -> Pin> + Send>> + Send + Sync { + let ud = make_user_data(user_data, free_fn); + move || { + let ud = Arc::clone(&ud); + Box::pin(async move { + clear_last_error(); + match unsafe { cb(ud.ptr) } { + NemoRelayStatus::Ok => Ok(()), + status => Err(FlowError::Internal(last_error_message().unwrap_or_else( + || format!("export target activation failed with status {status:?}"), + ))), + } + }) + } +} + // --------------------------------------------------------------------------- // Wrapper functions: C callback -> core trait objects // --------------------------------------------------------------------------- diff --git a/crates/node/observability.d.ts b/crates/node/observability.d.ts index e80f6b484..eebe0b260 100644 --- a/crates/node/observability.d.ts +++ b/crates/node/observability.d.ts @@ -6,6 +6,12 @@ import type { ConfigPolicy, ConfigDiagnostic, ConfigReport } from './plugin'; export { ConfigPolicy, ConfigDiagnostic, ConfigReport }; +export interface ExportActivationPolicyConfig { + provider: string; + timeout_millis?: number; + config?: Json; +} + export interface AtofConfig { enabled?: boolean; sinks?: AtofSinkConfig[]; @@ -18,6 +24,7 @@ export interface AtofFileSinkConfig { output_directory?: string; filename?: string; mode?: 'append' | 'overwrite' | string; + activation_policy?: ExportActivationPolicyConfig; } export interface AtofStreamSinkConfig { @@ -29,6 +36,7 @@ export interface AtofStreamSinkConfig { timeout_millis?: number; field_name_policy?: 'preserve' | 'replace_dots' | string; name?: string; + activation_policy?: ExportActivationPolicyConfig; } /** @deprecated Use AtofStreamSinkConfig. */ @@ -44,6 +52,7 @@ export interface S3StorageConfig { region?: string; endpoint_url?: string; allow_http?: boolean; + activation_policy?: ExportActivationPolicyConfig; } export interface HttpStorageConfig { @@ -52,6 +61,7 @@ export interface HttpStorageConfig { headers?: Record; header_env?: Record; timeout_millis?: number; + activation_policy?: ExportActivationPolicyConfig; } export interface AtifConfig { @@ -64,6 +74,7 @@ export interface AtifConfig { output_directory?: string; filename_template?: string; storage?: S3StorageConfig | HttpStorageConfig | Array; + local_activation_policy?: ExportActivationPolicyConfig; } export interface OpenTelemetryEndpointConfig { @@ -85,6 +96,7 @@ export interface OpenTelemetryEndpointConfig { max_queue_size?: number; max_export_batch_size?: number; scheduled_delay_millis?: number; + activation_policy?: ExportActivationPolicyConfig; } export interface OpenTelemetrySignalEndpointConfig { @@ -98,6 +110,7 @@ export interface OpenTelemetrySignalEndpointConfig { service_version?: string; instrumentation_scope?: string; timeout_millis?: number; + activation_policy?: ExportActivationPolicyConfig; } export interface OpenTelemetryLogSectionConfig { @@ -149,6 +162,10 @@ export declare function defaultConfig(): Config; export declare function atofConfig(config?: AtofConfig): AtofConfig; /** Create per-agent Agent Trajectory Interchange Format (ATIF) trajectory settings with defaults applied. */ export declare function atifConfig(config?: AtifConfig): AtifConfig; +/** Create and validate one export-activation policy attachment. */ +export declare function exportActivationPolicy( + config: ExportActivationPolicyConfig, +): ExportActivationPolicyConfig; /** Create one typed OpenTelemetry endpoint. */ export declare function openTelemetryEndpoint(config: OpenTelemetryEndpointConfig): OpenTelemetryEndpointConfig; /** Create one signal-specific OpenTelemetry endpoint for logs or metrics. */ diff --git a/crates/node/observability.js b/crates/node/observability.js index c4bfc6620..4fe93d5d8 100644 --- a/crates/node/observability.js +++ b/crates/node/observability.js @@ -47,6 +47,30 @@ function atifConfig(config = {}) { }; } +/** + * Create and validate one export-activation policy attachment. + * + * @param {object} config - Policy settings including a dynamic plugin provider ID. + * @returns {object} A normalized policy attachment. + */ +function exportActivationPolicy(config) { + if (!config || typeof config !== 'object') { + throw new TypeError('Export activation policy config is required'); + } + if (typeof config.provider !== 'string' || config.provider.trim() === '') { + throw new TypeError('Export activation policy provider must be a nonblank string'); + } + const timeout = config.timeout_millis ?? 30000; + if (!Number.isInteger(timeout) || timeout < 1000 || timeout > 300000) { + throw new TypeError('Export activation policy timeout_millis must be between 1000 and 300000'); + } + return { + ...config, + timeout_millis: timeout, + config: config.config ?? null, + }; +} + /** * Create one typed OpenTelemetry endpoint. * @@ -167,6 +191,7 @@ module.exports = { defaultConfig, atofConfig, atifConfig, + exportActivationPolicy, openTelemetryEndpoint, openTelemetrySignalEndpoint, openTelemetryLogConfig, diff --git a/crates/node/plugin.d.ts b/crates/node/plugin.d.ts index 8f1d689e0..a66e0f6c2 100644 --- a/crates/node/plugin.d.ts +++ b/crates/node/plugin.d.ts @@ -204,8 +204,37 @@ export interface ToolExecutionInterceptOutcome { pendingMarks?: PendingMarkSpec[]; } +/** Configuration selecting a component-scoped export activation provider. */ +export interface ExportActivationPolicyConfig { + provider: string; + timeout_millis?: number; + config?: Json; +} + +/** Request passed to an export activation policy exactly once per target. */ +export interface ExportActivationRequest { + target_kind: string; + config: Json; +} + +/** Deferred registration for one plugin-managed export target. */ +export interface ExportTargetRegistration { + id: string; + targetKind: string; + activationPolicy?: ExportActivationPolicyConfig; +} + /** Component-scoped registration context passed to plugin handlers. */ export interface PluginContext { + /** Register this component's export activation policy provider. */ + registerExportActivationPolicy( + callback: (request: ExportActivationRequest) => "allow" | "deny" | Promise<"allow" | "deny">, + ): void; + /** Register a target whose callback constructs and starts the exporter only when allowed. */ + registerExportTarget( + registration: ExportTargetRegistration, + activate: () => void | Promise, + ): void; /** * Register an event subscriber for this component. Callback failures are isolated and reported * through the Node binding's callback-error channel; flushSubscribers waits for returned promises. diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index 0935a3f40..897dc0f97 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -28,6 +28,10 @@ use serde::Deserialize; use serde_json::Value as Json; use tokio_stream::{Stream, StreamExt}; +use nemo_relay::api::export_activation::{ + ExportActivationDecision, ExportActivationPolicyConfig, ExportActivationPolicyRegistry, + ExportActivationRequest, ExportActivationTargetKind, ExportTargetRegistration, +}; use nemo_relay::api::llm as core_llm_api; use nemo_relay::api::llm::{LlmAttributes, LlmRequest}; use nemo_relay::api::registry as core_registry_api; @@ -823,6 +827,8 @@ fn add_plugin_event_sanitizer( fn build_plugin_context( env: &Env, namespace_prefix: String, + provider_id: String, + export_activation: Arc, registrations: Arc>>, ) -> napi::Result { let mut context = env.create_object()?; @@ -1332,7 +1338,7 @@ fn build_plugin_context( )?; let tool_regs = registrations.clone(); - let tool_exec_namespace = namespace_prefix; + let tool_exec_namespace = namespace_prefix.clone(); let register_tool_execution_intercept = env.create_function_from_closure( "__nemo_relay_adaptive_register_tool_execution_intercept", move |ctx| { @@ -1370,12 +1376,129 @@ fn build_plugin_context( register_tool_execution_intercept, )?; + let policy_registry = export_activation.clone(); + let policy_provider_id = provider_id.clone(); + let policy_regs = registrations.clone(); + let register_export_activation_policy = env.create_function_from_closure( + "__nemo_relay_plugin_register_export_activation_policy", + move |ctx| { + let callback = Arc::new(PromiseAwareFn::new(ctx.env, &ctx.get::(0)?)?); + let policy_callback = callback.clone(); + policy_registry + .register( + &policy_provider_id, + Arc::new(move |request: ExportActivationRequest| { + let policy_callback = policy_callback.clone(); + Box::pin(async move { + let request = serde_json::to_value(request).map_err(|error| { + FlowError::Internal(format!( + "failed to serialize export activation request: {error}" + )) + })?; + let decision = policy_callback.call(request).await?; + serde_json::from_value::(decision).map_err( + |error| { + FlowError::Internal(format!( + "invalid export activation decision: {error}" + )) + }, + ) + }) + }), + ) + .map_err(to_napi_err)?; + let cleanup_registry = policy_registry.clone(); + let cleanup_provider_id = policy_provider_id.clone(); + policy_regs.lock().unwrap().push(PluginRegistration::new( + "export_activation_policy", + cleanup_provider_id.clone(), + Box::new(move || { + cleanup_registry + .deregister(&cleanup_provider_id) + .map(|_| ()) + .map_err(|error| { + PluginError::RegistrationFailed(format!( + "export activation policy deregistration failed: {error}" + )) + })?; + Ok(()) + }), + )); + ctx.env.get_undefined() + }, + )?; + context.set_named_property( + "registerExportActivationPolicy", + register_export_activation_policy, + )?; + + #[derive(Deserialize)] + #[serde(rename_all = "camelCase", deny_unknown_fields)] + struct NodeExportTargetRegistration { + id: String, + target_kind: String, + #[serde(default)] + activation_policy: Option, + } + + let target_registry = export_activation; + let target_regs = registrations; + let register_export_target = env.create_function_from_closure( + "__nemo_relay_plugin_register_export_target", + move |ctx| { + let registration = ctx.get::(0)?; + let registration: NodeExportTargetRegistration = serde_json::from_value(registration) + .map_err(|error| { + napi::Error::from_reason(format!("invalid export target: {error}")) + })?; + let target_kind = ExportActivationTargetKind::new(registration.target_kind) + .map_err(napi::Error::from_reason)?; + let target_id = registration.id.clone(); + let qualified_target_id = format!("{namespace_prefix}export-target:{target_id}"); + let callback = Arc::new(PromiseAwareFn::new(ctx.env, &ctx.get::(1)?)?); + target_registry + .register_target( + qualified_target_id.clone(), + ExportTargetRegistration { + id: registration.id, + target_kind, + activation_policy: registration.activation_policy, + }, + Arc::new(move || { + let callback = callback.clone(); + Box::pin(async move { callback.call_spread(Vec::new()).await.map(|_| ()) }) + }), + ) + .map_err(to_napi_err)?; + let cleanup_registry = target_registry.clone(); + let cleanup_target_id = qualified_target_id; + target_regs.lock().unwrap().push(PluginRegistration::new( + "export_target", + target_id, + Box::new(move || { + cleanup_registry + .deregister_target(&cleanup_target_id) + .map(|_| ()) + .map_err(|error| { + PluginError::RegistrationFailed(format!( + "export target deregistration failed: {error}" + )) + }) + }), + )); + ctx.env.get_undefined() + }, + )?; + context.set_named_property("registerExportTarget", register_export_target)?; + Ok(context) } struct NodePluginRegisterCall { plugin_config: Json, namespace_prefix: String, + provider_id: String, + export_activation: Arc, registrations: Arc>>, } @@ -1740,6 +1863,8 @@ impl Plugin for NodePlugin { let payload = NodePluginRegisterCall { plugin_config: Json::Object(plugin_config), namespace_prefix, + provider_id: self.plugin_kind.clone(), + export_activation: ctx.export_activation_policies().clone(), registrations: registrations.clone(), }; let (tx, rx) = std::sync::mpsc::sync_channel::>(1); @@ -5208,6 +5333,8 @@ pub fn register_plugin( let plugin_context = build_plugin_context( &ctx.env, ctx.value.namespace_prefix, + ctx.value.provider_id, + ctx.value.export_activation, ctx.value.registrations, )?; Ok(vec![ diff --git a/crates/node/tests/observability_plugin_tests.mjs b/crates/node/tests/observability_plugin_tests.mjs index 46079a33d..54aa689f3 100644 --- a/crates/node/tests/observability_plugin_tests.mjs +++ b/crates/node/tests/observability_plugin_tests.mjs @@ -18,6 +18,53 @@ function tempDir(prefix) { } describe('observability plugin helpers', () => { + it('builds export activation policies for remote targets', () => { + const policy = observability.exportActivationPolicy({ + provider: 'com.example.runtime-policy', + config: { allowed_countries: ['US', 'CA'] }, + }); + assert.deepEqual(policy, { + provider: 'com.example.runtime-policy', + timeout_millis: 30000, + config: { allowed_countries: ['US', 'CA'] }, + }); + assert.deepEqual( + observability.exportActivationPolicy({ + provider: 'com.example.runtime-policy', + timeout_millis: undefined, + config: undefined, + }), + { + provider: 'com.example.runtime-policy', + timeout_millis: 30000, + config: null, + }, + ); + assert.equal( + observability.openTelemetryEndpoint({ + type: 'full', + endpoint: 'http://localhost:4318/v1/traces', + activation_policy: policy, + }).activation_policy, + policy, + ); + assert.equal( + observability.atofConfig({ sinks: [{ type: 'file', activation_policy: policy }] }).sinks[0] + .activation_policy, + policy, + ); + assert.equal(observability.atifConfig({ local_activation_policy: policy }).local_activation_policy, policy); + assert.throws(() => observability.exportActivationPolicy({ provider: ' ' }), /nonblank/); + assert.throws( + () => observability.exportActivationPolicy({ provider: 'test', timeout_millis: 999 }), + /between 1000 and 300000/, + ); + assert.throws( + () => observability.exportActivationPolicy({ provider: 'test', timeout_millis: 300001 }), + /between 1000 and 300000/, + ); + }); + it('builds defaults and plugin component shape', () => { assert.deepEqual(observability.defaultConfig(), { version: 4 }); assert.equal( diff --git a/crates/node/tests/plugin_export_activation_tests.mjs b/crates/node/tests/plugin_export_activation_tests.mjs new file mode 100644 index 000000000..540220ec8 --- /dev/null +++ b/crates/node/tests/plugin_export_activation_tests.mjs @@ -0,0 +1,56 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, it } from 'node:test'; +import assert from 'node:assert/strict'; +import { createRequire } from 'node:module'; + +const require = createRequire(import.meta.url); +const plugin = require('../plugin.js'); + +describe('plugin export activation hooks', () => { + for (const [decision, expected] of [ + ['allow', 1], + ['deny', 0], + ]) { + it(`lets one plugin ${decision} its own exporter`, async () => { + const kind = `tests.node_export_activation_${decision}`; + let activations = 0; + plugin.register(kind, { + register(_config, context) { + context.registerExportActivationPolicy(async (request) => { + assert.deepEqual(request, { + target_kind: 'tests.telemetry.otlp', + config: { country: 'US' }, + }); + return decision; + }); + context.registerExportTarget( + { + id: 'self-otel', + targetKind: 'tests.telemetry.otlp', + activationPolicy: { + provider: kind, + timeout_millis: 30000, + config: { country: 'US' }, + }, + }, + async () => { + activations += 1; + }, + ); + }, + }); + try { + await plugin.initialize({ + version: 1, + components: [{ kind, enabled: true, config: {} }], + }); + assert.equal(activations, expected); + } finally { + plugin.clear(); + plugin.deregister(kind); + } + }); + } +}); diff --git a/crates/node/tests/public_observability_api_fixture.ts b/crates/node/tests/public_observability_api_fixture.ts index 6c1bd62e9..6471a91c3 100644 --- a/crates/node/tests/public_observability_api_fixture.ts +++ b/crates/node/tests/public_observability_api_fixture.ts @@ -12,6 +12,21 @@ import { OpenTelemetryLogSubscriber, OpenTelemetryMetricSubscriber, } from '../index.js'; +import { + type ExportActivationPolicyConfig, + exportActivationPolicy, + openTelemetryEndpoint, +} from '../observability.js'; + +const activationPolicy: ExportActivationPolicyConfig = exportActivationPolicy({ + provider: 'com.example.runtime-policy', + config: { allowed_countries: ['US', 'CA'] }, +}); +openTelemetryEndpoint({ + type: 'full', + endpoint: 'http://localhost:4318/v1/traces', + activation_policy: activationPolicy, +}); const dataSchema: DataSchema = { name: 'example.fixture', version: '1' }; diff --git a/crates/plugin/src/async_sdk.rs b/crates/plugin/src/async_sdk.rs index 43b636e98..cf615b290 100644 --- a/crates/plugin/src/async_sdk.rs +++ b/crates/plugin/src/async_sdk.rs @@ -1103,6 +1103,33 @@ impl PluginContext<'_> { } } + fn register_activation_adapter( + &mut self, + kind: NemoRelayNativeActivationHookKind, + payload: &str, + adapter: Box, + ) -> Result<()> { + let state = Box::into_raw(Box::new(UnaryCallbackState { + host: self.host_v4()?, + executor: Arc::clone(&self.executor), + adapter, + })); + let status = unsafe { + self.register_activation_hook_raw( + kind, + payload, + unary_trampoline, + state.cast(), + Some(drop_unary_callback), + ) + }; + if status == NemoRelayStatus::Ok { + Ok(()) + } else { + Err(status_message(self.host, status, "activation hook")) + } + } + fn register_event_adapter( &mut self, kind: NemoRelayNativeAsyncMiddlewareKind, @@ -1165,6 +1192,56 @@ impl PluginContext<'_> { ) } + /// Registers this plugin's activation-time export policy. + pub fn register_export_activation_policy(&mut self, callback: F) -> Result<()> + where + F: Fn(ExportActivationRequest) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let callback = Arc::new(callback); + self.register_activation_adapter( + NemoRelayNativeActivationHookKind::ExportActivationPolicy, + "export_activation_policy", + Box::new(move |value, _, _| { + let callback = Arc::clone(&callback); + Box::pin(async move { + let request: ExportActivationRequest = + serde_json::from_value(value).map_err(|error| { + format!("invalid export activation policy request: {error}") + })?; + serde_json::to_value(callback(request).await?) + .map_err(|error| error.to_string()) + }) + }), + ) + } + + /// Registers a deferred local or remote export target. + pub fn register_export_target( + &mut self, + registration: ExportTargetRegistration, + activate: F, + ) -> Result<()> + where + F: Fn() -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let payload = serde_json::to_string(®istration) + .map_err(|error| format!("failed to serialize export target registration: {error}"))?; + let activate = Arc::new(activate); + self.register_activation_adapter( + NemoRelayNativeActivationHookKind::ExportTarget, + &payload, + Box::new(move |_, _, _| { + let activate = Arc::clone(&activate); + Box::pin(async move { + activate().await?; + Ok(Json::Null) + }) + }), + ) + } + /// Registers an asynchronous mark-event sanitizer. pub fn register_mark_sanitize_guardrail( &mut self, diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 5028bc9ec..8873889a4 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -40,7 +40,10 @@ pub use nemo_relay_types::codec::optimization::{ }; pub use nemo_relay_types::codec::request::AnnotatedLlmRequest; pub use nemo_relay_types::codec::response::AnnotatedLlmResponse; -pub use nemo_relay_types::plugin::{ConfigDiagnostic, DiagnosticLevel}; +pub use nemo_relay_types::plugin::{ + ConfigDiagnostic, DiagnosticLevel, ExportActivationDecision, ExportActivationPolicyConfig, + ExportActivationRequest, ExportActivationTargetKind, ExportTargetRegistration, +}; use serde::{Serialize, de::DeserializeOwned}; use serde_json::Map; @@ -855,6 +858,28 @@ impl TryFrom for NemoRelayNativeAsyncMiddlewareKind { } } +/// Activation-hook surface selected by the native generic async transport. +#[repr(u32)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum NemoRelayNativeActivationHookKind { + /// Policy provider evaluated before export-target construction. + ExportActivationPolicy = 15, + /// Deferred local or remote export-target constructor. + ExportTarget = 16, +} + +impl TryFrom for NemoRelayNativeActivationHookKind { + type Error = (); + + fn try_from(value: u32) -> std::result::Result { + match value { + 15 => Ok(Self::ExportActivationPolicy), + 16 => Ok(Self::ExportTarget), + _ => Err(()), + } + } +} + /// Indicates whether an asynchronous native callback settled before returning. #[repr(u32)] #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -2474,6 +2499,46 @@ impl<'a> PluginContext<'a> { }) } + /// Registers an activation hook through the ABI-v4 generic async transport. + /// + /// This does not add fields to the ABI table; activation hooks use distinct + /// discriminants carried by the existing completion-based callback transport. + /// + /// # Safety + /// `cb`, `user_data`, and `free_fn` must remain valid until the host + /// deregisters the callback or invokes `free_fn`. + #[allow(clippy::too_many_arguments)] + pub unsafe fn register_activation_hook_raw( + &mut self, + kind: NemoRelayNativeActivationHookKind, + payload: &str, + cb: NemoRelayNativeAsyncMiddlewareCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus { + if self.host.abi_version < NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE + || self.host.struct_size < std::mem::size_of::() + { + if let Some(free_fn) = free_fn { + unsafe { free_fn(user_data) }; + } + return NemoRelayStatus::InvalidArg; + } + let host = unsafe { &*(self.host as *const _ as *const NemoRelayNativeHostApiV3) }; + self.with_name_and_callback(payload, user_data, free_fn, |_, payload| unsafe { + (host.plugin_context_register_async_middleware)( + self.raw, + kind as u32, + payload, + 0, + false, + cb, + user_data, + free_fn, + ) + }) + } + /// Registers an incremental completion-based LLM stream intercept. /// /// # Safety diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 6489a724f..846c01f50 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -22,15 +22,17 @@ use std::time::{Duration, Instant}; use futures::StreamExt; use nemo_relay_plugin::{ AnnotatedLlmRequest, BuiltinLlmCodec, CategoryProfile, ConfigDiagnostic, DataSchema, - DiagnosticLevel, Event, EventCategory, EventSanitizeFields, Json, LlmCodecIdentity, + DiagnosticLevel, Event, EventCategory, EventSanitizeFields, ExportActivationDecision, + ExportActivationTargetKind, ExportTargetRegistration, Json, LlmCodecIdentity, LlmJsonAsyncStream, LlmJsonStream, LlmNext, LlmRequest, LlmRequestInterceptOutcome, LlmStream, LlmStreamNext, LogSeverity, MetricKind, MetricMeasurement, MetricValueType, NEMO_RELAY_NATIVE_ABI_VERSION, NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE, NEMO_RELAY_NATIVE_ABI_VERSION_LEGACY, NativeExecutorConfig, NativePlugin, - NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, - NemoRelayNativeAsyncLlmStreamOpenCb, NemoRelayNativeAsyncLlmStreamPullCb, - NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, - NemoRelayNativeAsyncNextResultCb, NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, + NemoRelayNativeActivationHookKind, NemoRelayNativeAsyncCallbackState, + NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncLlmStreamOpenCb, + NemoRelayNativeAsyncLlmStreamPullCb, NemoRelayNativeAsyncMiddlewareCb, + NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncNextResultCb, + NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativeLlmAsyncStream, @@ -74,6 +76,11 @@ fn async_abi_discriminants_reject_unknown_values() { assert_eq!(Kind::try_from(discriminant as u32), Ok(kind)); } assert!(NemoRelayNativeAsyncMiddlewareKind::try_from(15).is_err()); + assert_eq!( + NemoRelayNativeActivationHookKind::ExportActivationPolicy as u32, + 15 + ); + assert_eq!(NemoRelayNativeActivationHookKind::ExportTarget as u32, 16); assert_eq!( NemoRelayNativeAsyncCallbackState::try_from(1), Ok(NemoRelayNativeAsyncCallbackState::Pending) @@ -286,7 +293,7 @@ struct RegisteredLlmRequestIntercept { } struct RegisteredAsync { - kind: NemoRelayNativeAsyncMiddlewareKind, + kind: u32, name: String, priority: i32, break_chain: bool, @@ -1854,10 +1861,6 @@ unsafe extern "C" fn capture_register_async_middleware( } return status; } - let kind = match NemoRelayNativeAsyncMiddlewareKind::try_from(kind) { - Ok(kind) => kind, - Err(()) => return NemoRelayStatus::InvalidArg, - }; let name = match required_host_string(&test_host(), name) { Ok(name) => name, Err(status) => return status, @@ -2167,11 +2170,19 @@ fn test_llm_request() -> LlmRequest { } fn take_async_registration(kind: NemoRelayNativeAsyncMiddlewareKind) -> RegisteredAsync { + take_async_registration_raw(kind as u32) +} + +fn take_activation_registration(kind: NemoRelayNativeActivationHookKind) -> RegisteredAsync { + take_async_registration_raw(kind as u32) +} + +fn take_async_registration_raw(kind: u32) -> RegisteredAsync { let mut registrations = ASYNC_REGISTRATIONS.lock().unwrap(); let index = registrations .iter() .position(|registration| registration.kind == kind) - .unwrap_or_else(|| panic!("missing {kind:?} registration")); + .unwrap_or_else(|| panic!("missing kind {kind} registration")); registrations.remove(index) } @@ -2182,7 +2193,7 @@ fn invoke_async_registration( next: Option<&MockAsyncNext>, ) -> std::result::Result { ASYNC_TOOL_NEXT_RESULT.store( - registration.kind == NemoRelayNativeAsyncMiddlewareKind::ToolExecutionIntercept, + registration.kind == NemoRelayNativeAsyncMiddlewareKind::ToolExecutionIntercept as u32, Ordering::SeqCst, ); let completion = MockAsyncCompletion::new(); @@ -3295,6 +3306,21 @@ fn typed_async_middleware_registers_and_round_trips_every_surface() { )])) }) .unwrap(); + ctx.register_export_activation_policy(|request| async move { + assert_eq!(request.target_kind, ExportActivationTargetKind::OtlpTrace); + assert_eq!(request.config, json!({"enabled": true})); + Ok(ExportActivationDecision::Allow) + }) + .unwrap(); + ctx.register_export_target( + ExportTargetRegistration { + id: "self-otel".into(), + target_kind: ExportActivationTargetKind::OTLP_TRACE, + activation_policy: None, + }, + || async { Ok(()) }, + ) + .unwrap(); let metadata = ASYNC_REGISTRATIONS .lock() @@ -3309,17 +3335,17 @@ fn typed_async_middleware_registers_and_round_trips_every_surface() { ) }) .collect::>(); - assert_eq!(metadata.len(), 14); + assert_eq!(metadata.len(), 16); assert_eq!(metadata[0].1, "mark-async"); assert_eq!(metadata[0].2, 1); assert_eq!( metadata[6].0, - NemoRelayNativeAsyncMiddlewareKind::ToolRequestIntercept + NemoRelayNativeAsyncMiddlewareKind::ToolRequestIntercept as u32 ); assert!(metadata[6].3); assert_eq!( metadata[11].0, - NemoRelayNativeAsyncMiddlewareKind::LlmRequestIntercept + NemoRelayNativeAsyncMiddlewareKind::LlmRequestIntercept as u32 ); assert!(metadata[11].3); { @@ -3353,6 +3379,26 @@ fn typed_async_middleware_registers_and_round_trips_every_surface() { assert_eq!(result, json!({ "plugin.event_name": "checkpoint" })); unsafe { metadata_registration.free() }; + let policy_registration = + take_activation_registration(NemoRelayNativeActivationHookKind::ExportActivationPolicy); + let result = invoke_async_registration( + &host, + &policy_registration, + json!({"target_kind": "nemo_relay.otlp.trace", "config": {"enabled": true}}), + None, + ) + .unwrap(); + assert_eq!(result, json!("allow")); + unsafe { policy_registration.free() }; + + let target_registration = + take_activation_registration(NemoRelayNativeActivationHookKind::ExportTarget); + let target: ExportTargetRegistration = serde_json::from_str(&target_registration.name).unwrap(); + assert_eq!(target.id, "self-otel"); + let result = invoke_async_registration(&host, &target_registration, Json::Null, None).unwrap(); + assert_eq!(result, Json::Null); + unsafe { target_registration.free() }; + for kind in [ NemoRelayNativeAsyncMiddlewareKind::MarkSanitize, NemoRelayNativeAsyncMiddlewareKind::ScopeSanitizeStart, @@ -3548,6 +3594,47 @@ fn typed_async_middleware_registers_and_round_trips_every_surface() { assert_eq!(live_host_strings(), 0); } +#[test] +fn typed_export_activation_policy_round_trips_deny_and_error() { + let _guard = begin_test(); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + + ctx.register_export_activation_policy(|request| async move { + assert_eq!(request.target_kind, ExportActivationTargetKind::AtifS3); + assert_eq!(request.config, json!({"enabled": false})); + Ok(ExportActivationDecision::Deny) + }) + .unwrap(); + ctx.register_export_activation_policy(|_| async move { Err("policy failed".into()) }) + .unwrap(); + + let deny = + take_activation_registration(NemoRelayNativeActivationHookKind::ExportActivationPolicy); + let result = invoke_async_registration( + &host, + &deny, + json!({"target_kind": "nemo_relay.atif.s3", "config": {"enabled": false}}), + None, + ) + .unwrap(); + assert_eq!(result, json!("deny")); + unsafe { deny.free() }; + + let failure = + take_activation_registration(NemoRelayNativeActivationHookKind::ExportActivationPolicy); + let error = invoke_async_registration( + &host, + &failure, + json!({"target_kind": "nemo_relay.otlp.trace", "config": null}), + None, + ) + .unwrap_err(); + assert_eq!(error, "policy failed"); + unsafe { failure.free() }; + assert_eq!(live_host_strings(), 0); +} + #[test] fn typed_async_llm_sanitize_context_decodes_oci_genai_builtin_identity() { let _guard = begin_test(); diff --git a/crates/python/src/py_callable.rs b/crates/python/src/py_callable.rs index 266238f47..e8a66c265 100644 --- a/crates/python/src/py_callable.rs +++ b/crates/python/src/py_callable.rs @@ -46,6 +46,9 @@ use tokio_stream::Stream; use tokio_stream::wrappers::ReceiverStream; use nemo_relay::api::event::{Event, EventSanitizeFields}; +use nemo_relay::api::export_activation::{ + ExportActivationDecision, ExportActivationPolicyFn, ExportActivationRequest, +}; use nemo_relay::api::llm::LlmRequest; use nemo_relay::api::tool::ToolExecutionResult; use nemo_relay::codec::request::AnnotatedLlmRequest as AnnotatedLLMRequest; @@ -879,6 +882,74 @@ pub fn wrap_py_tool_fn(py_fn: Py) -> ToolSanitizeFn { }) } +/// Wraps a Python export activation policy callback. +pub fn wrap_py_export_activation_policy_fn(py_fn: Py) -> ExportActivationPolicyFn { + let py_fn = Arc::new(py_fn); + let task_locals = capture_python_task_locals(); + Arc::new(move |request: ExportActivationRequest| { + let py_fn = Arc::clone(&py_fn); + let task_locals = task_locals_with_running_loop(task_locals.as_ref()); + Box::pin(async move { + let result = resolve_py_object_or_future(Python::attach(|py| { + let request = json_to_py( + py, + &serde_json::to_value(request).map_err(|error| { + FlowError::Internal(format!( + "failed to serialize export activation request: {error}" + )) + })?, + ) + .map_err(|error| FlowError::Internal(error.to_string()))?; + let result = py_fn + .bind(py) + .call1((request,)) + .map_err(|error| FlowError::Internal(error.to_string()))?; + split_py_object_or_future_with_locals( + py, + result.unbind(), + task_locals.as_ref(), + None, + ) + })) + .await?; + Python::attach(|py| { + let value = py_to_json(result.bind(py)) + .map_err(|error| FlowError::Internal(error.to_string()))?; + serde_json::from_value::(value) + .map_err(|error| FlowError::Internal(error.to_string())) + }) + }) + }) +} + +/// Wraps a one-shot Python export-target activation callback. +pub fn wrap_py_export_target_activation_fn( + py_fn: Py, +) -> impl Fn() -> Pin> + Send>> + Send + Sync { + let py_fn = Arc::new(py_fn); + let task_locals = capture_python_task_locals(); + move || { + let py_fn = Arc::clone(&py_fn); + let task_locals = task_locals.clone(); + Box::pin(async move { + let _ = resolve_py_object_or_future(Python::attach(|py| { + let result = py_fn + .bind(py) + .call0() + .map_err(|error| FlowError::Internal(error.to_string()))?; + split_py_object_or_future_with_locals( + py, + result.unbind(), + task_locals.as_ref(), + None, + ) + })) + .await?; + Ok(()) + }) + } +} + /// Wrap a Python callable `(str, Json) -> Optional[str]` for tool conditional guardrails. pub fn wrap_py_tool_conditional_fn(py_fn: Py) -> ToolConditionalFn { let py_fn = Arc::new(py_fn); diff --git a/crates/python/src/py_plugin.rs b/crates/python/src/py_plugin.rs index ec11ce32e..2e4eb7752 100644 --- a/crates/python/src/py_plugin.rs +++ b/crates/python/src/py_plugin.rs @@ -12,6 +12,9 @@ use std::sync::{Arc, LazyLock, Mutex}; use pyo3::prelude::*; use serde_json::{Map, Value as Json}; +use nemo_relay::api::export_activation::{ + ExportActivationPolicyRegistry, ExportTargetRegistration, +}; use nemo_relay::api::registry::{ deregister_llm_conditional_execution_guardrail, deregister_llm_execution_intercept, deregister_llm_request_intercept, deregister_llm_sanitize_request_guardrail, @@ -39,11 +42,12 @@ use nemo_relay::plugin::{ use crate::convert::{json_to_py, py_to_json}; use crate::py_callable::{ - wrap_py_event_sanitize_fn, wrap_py_event_subscriber, wrap_py_llm_conditional_fn, - wrap_py_llm_exec_intercept_fn, wrap_py_llm_request_intercept_fn, - wrap_py_llm_sanitize_request_fn, wrap_py_llm_sanitize_response_fn, - wrap_py_llm_stream_exec_intercept_fn, wrap_py_tool_conditional_fn, - wrap_py_tool_exec_intercept_fn, wrap_py_tool_fn, wrap_py_tool_request_intercept_fn, + wrap_py_event_sanitize_fn, wrap_py_event_subscriber, wrap_py_export_activation_policy_fn, + wrap_py_export_target_activation_fn, wrap_py_llm_conditional_fn, wrap_py_llm_exec_intercept_fn, + wrap_py_llm_request_intercept_fn, wrap_py_llm_sanitize_request_fn, + wrap_py_llm_sanitize_response_fn, wrap_py_llm_stream_exec_intercept_fn, + wrap_py_tool_conditional_fn, wrap_py_tool_exec_intercept_fn, wrap_py_tool_fn, + wrap_py_tool_request_intercept_fn, }; #[cfg(test)] @@ -144,6 +148,8 @@ fn new_py_plugin_context( _plugin_kind: &str, registrations: Arc>>, namespace_prefix: String, + provider_id: String, + export_activation: Arc, ) -> PyResult> { #[cfg(test)] if FORCE_PLUGIN_CONTEXT_NEW_ERROR @@ -161,6 +167,8 @@ fn new_py_plugin_context( PyPluginContext { registrations, namespace_prefix, + provider_id, + export_activation, }, ) } @@ -171,12 +179,16 @@ pub(crate) fn invoke_python_plugin_register( register_fn: &Bound<'_, PyAny>, plugin_config: &Map, namespace_prefix: String, + provider_id: String, + export_activation: Arc, ) -> PyResult> { let py_ctx = new_py_plugin_context( py, plugin_kind, Arc::new(Mutex::new(vec![])), namespace_prefix, + provider_id, + export_activation, )?; let plugin_config_py = plugin_config_to_py(py, plugin_kind, plugin_config)?; match register_fn.call1((plugin_config_py, py_ctx.clone_ref(py))) { @@ -197,6 +209,8 @@ pub(crate) fn invoke_python_plugin_register( pub struct PyPluginContext { registrations: Arc>>, namespace_prefix: String, + provider_id: String, + export_activation: Arc, } impl PyPluginContext { @@ -240,6 +254,74 @@ impl PyPluginContext { #[pymethods] impl PyPluginContext { + #[pyo3(signature = (callback: "object") -> "None", text_signature = "(callback: object) -> None")] + fn register_export_activation_policy(&self, callback: Py) -> PyResult<()> { + self.export_activation + .register( + &self.provider_id, + wrap_py_export_activation_policy_fn(callback), + ) + .map_err(to_py_err)?; + let provider = self.provider_id.clone(); + let export_activation = Arc::clone(&self.export_activation); + self.registrations + .lock() + .map_err(|error| { + pyo3::exceptions::PyRuntimeError::new_err(format!( + "plugin context lock poisoned: {error}" + )) + })? + .push(PluginRegistration::new( + "export_activation_policy", + provider.clone(), + Box::new(move || { + export_activation + .deregister(&provider) + .map(|_| ()) + .map_err(|error| PluginError::RegistrationFailed(error.to_string())) + }), + )); + Ok(()) + } + + #[pyo3(signature = (registration: "object", callback: "object") -> "None", text_signature = "(registration: object, callback: object) -> None")] + fn register_export_target( + &self, + registration: &Bound<'_, PyAny>, + callback: Py, + ) -> PyResult<()> { + let registration = + serde_json::from_value::(py_to_json(registration)?) + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?; + let qualified_id = self.qualify_name(&format!("export-target:{}", registration.id)); + self.export_activation + .register_target( + qualified_id.clone(), + registration, + Arc::new(wrap_py_export_target_activation_fn(callback)), + ) + .map_err(to_py_err)?; + let export_activation = Arc::clone(&self.export_activation); + self.registrations + .lock() + .map_err(|error| { + pyo3::exceptions::PyRuntimeError::new_err(format!( + "plugin context lock poisoned: {error}" + )) + })? + .push(PluginRegistration::new( + "export_target", + qualified_id.clone(), + Box::new(move || { + export_activation + .deregister_target(&qualified_id) + .map(|_| ()) + .map_err(|error| PluginError::RegistrationFailed(error.to_string())) + }), + )); + Ok(()) + } + #[pyo3(signature = (name: "str", priority: "int", callback: "object") -> "None", text_signature = "(name: str, priority: int, callback: object) -> None")] fn register_mark_sanitize_guardrail( &self, @@ -642,6 +724,8 @@ impl Plugin for PyPlugin { ctx: &'a mut PluginRegistrationContext, ) -> Pin> + Send + 'a>> { let namespace_prefix = ctx.qualify_name(""); + let export_activation = ctx.export_activation_policies(); + let provider_id = self.plugin_kind.clone(); let plugin_config = plugin_config.clone(); Box::pin(async move { let registrations = Python::attach(|py| -> PyResult> { @@ -652,6 +736,8 @@ impl Plugin for PyPlugin { ®ister_fn, &plugin_config, namespace_prefix, + provider_id, + export_activation, ) }) .map_err(|err| PluginError::RegistrationFailed(err.to_string()))?; diff --git a/crates/python/tests/coverage/py_plugin_coverage_tests.rs b/crates/python/tests/coverage/py_plugin_coverage_tests.rs index e3973e587..ef30e68cb 100644 --- a/crates/python/tests/coverage/py_plugin_coverage_tests.rs +++ b/crates/python/tests/coverage/py_plugin_coverage_tests.rs @@ -53,6 +53,8 @@ fn plugin_context_helpers_and_error_conversion_work() { let context = PyPluginContext { registrations: Arc::new(Mutex::new(vec![])), namespace_prefix: "demo.".to_string(), + provider_id: "demo".to_string(), + export_activation: Arc::new(Default::default()), }; assert_eq!(context.qualify_name("subscriber"), "demo.subscriber"); @@ -73,6 +75,8 @@ fn plugin_context_rejects_legacy_and_uninspectable_llm_sanitizers() { let context = PyPluginContext { registrations: Arc::new(Mutex::new(vec![])), namespace_prefix: "invalid.".to_string(), + provider_id: "invalid".to_string(), + export_activation: Arc::new(Default::default()), }; Python::attach(|py| { @@ -263,6 +267,8 @@ async def tool_execution_intercept(name, value, next): let context = PyPluginContext { registrations: Arc::new(Mutex::new(vec![])), namespace_prefix: "demo.".to_string(), + provider_id: "demo".to_string(), + export_activation: Arc::new(Default::default()), }; context @@ -744,6 +750,8 @@ async def tool_execution_intercept(name, value, next): let context = PyPluginContext { registrations: Arc::new(Mutex::new(vec![])), namespace_prefix: "rollback.".to_string(), + provider_id: "rollback".to_string(), + export_activation: Arc::new(Default::default()), }; context @@ -966,6 +974,8 @@ class FailingPlugin: ®ister_fn, &serde_json::Map::new(), namespace_prefix.clone(), + "demo.rollback".to_string(), + Arc::new(Default::default()), ) .unwrap_err(); assert!(err.to_string().contains("boom"), "{err}"); @@ -973,6 +983,8 @@ class FailingPlugin: let context = PyPluginContext { registrations: Arc::new(Mutex::new(vec![])), namespace_prefix: namespace_prefix.clone(), + provider_id: "demo.rollback".to_string(), + export_activation: Arc::new(Default::default()), }; context .register_subscriber("sub", helpers.getattr("subscriber").unwrap().unbind()) @@ -1037,6 +1049,8 @@ async def tool_execution_intercept(name, value, next): let context = PyPluginContext { registrations, namespace_prefix: "poison.".to_string(), + provider_id: "poison".to_string(), + export_activation: Arc::new(Default::default()), }; fn assert_poisoned_tool_registrations( diff --git a/crates/types/src/plugin.rs b/crates/types/src/plugin.rs index eeb3be024..22ed4d7c7 100644 --- a/crates/types/src/plugin.rs +++ b/crates/types/src/plugin.rs @@ -3,7 +3,157 @@ //! Shared plugin diagnostic data types. -use serde::{Deserialize, Serialize}; +use std::borrow::Cow; + +use serde::{Deserialize, Deserializer, Serialize}; + +use crate::Json; + +/// Export target kind presented to an activation policy. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(transparent)] +pub struct ExportActivationTargetKind(Cow<'static, str>); + +impl ExportActivationTargetKind { + /// OpenTelemetry trace destination. + pub const OTLP_TRACE: Self = Self(Cow::Borrowed("nemo_relay.otlp.trace")); + /// OpenTelemetry log destination. + pub const OTLP_LOG: Self = Self(Cow::Borrowed("nemo_relay.otlp.log")); + /// OpenTelemetry metric destination. + pub const OTLP_METRIC: Self = Self(Cow::Borrowed("nemo_relay.otlp.metric")); + /// Local ATOF file sink. + pub const ATOF_FILE: Self = Self(Cow::Borrowed("nemo_relay.atof.file")); + /// Remote ATOF stream sink. + pub const ATOF_STREAM: Self = Self(Cow::Borrowed("nemo_relay.atof.stream")); + /// Local ATIF file destination. + pub const ATIF_FILE: Self = Self(Cow::Borrowed("nemo_relay.atif.file")); + /// Remote ATIF HTTP storage. + pub const ATIF_HTTP: Self = Self(Cow::Borrowed("nemo_relay.atif.http")); + /// Remote ATIF S3-compatible storage. + pub const ATIF_S3: Self = Self(Cow::Borrowed("nemo_relay.atif.s3")); + + // Compatibility spellings retained for the 0.8 pre-release SDK surface. + /// Compatibility alias for [`Self::OTLP_TRACE`]. + #[allow(non_upper_case_globals)] + pub const OtlpTrace: Self = Self::OTLP_TRACE; + /// Compatibility alias for [`Self::OTLP_LOG`]. + #[allow(non_upper_case_globals)] + pub const OtlpLog: Self = Self::OTLP_LOG; + /// Compatibility alias for [`Self::OTLP_METRIC`]. + #[allow(non_upper_case_globals)] + pub const OtlpMetric: Self = Self::OTLP_METRIC; + /// Compatibility alias for [`Self::ATOF_FILE`]. + #[allow(non_upper_case_globals)] + pub const AtofFile: Self = Self::ATOF_FILE; + /// Compatibility alias for [`Self::ATOF_STREAM`]. + #[allow(non_upper_case_globals)] + pub const AtofStream: Self = Self::ATOF_STREAM; + /// Compatibility alias for [`Self::ATIF_FILE`]. + #[allow(non_upper_case_globals)] + pub const AtifFile: Self = Self::ATIF_FILE; + /// Compatibility alias for [`Self::ATIF_HTTP`]. + #[allow(non_upper_case_globals)] + pub const AtifHttp: Self = Self::ATIF_HTTP; + /// Compatibility alias for [`Self::ATIF_S3`]. + #[allow(non_upper_case_globals)] + pub const AtifS3: Self = Self::ATIF_S3; + + /// Creates a custom namespaced export-target kind. + pub fn new(value: impl Into) -> Result { + let value = value.into(); + let valid = value.len() <= 255 + && value.contains('.') + && value.split('.').all(|segment| { + !segment.is_empty() + && segment + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-')) + }); + if !valid { + return Err( + "export activation target kind must be a dot-separated namespaced identifier" + .into(), + ); + } + Ok(Self(Cow::Owned(value))) + } + + /// Returns the serialized target-kind name. + pub fn as_str(&self) -> &str { + self.0.as_ref() + } +} + +impl<'de> Deserialize<'de> for ExportActivationTargetKind { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + Self::new(value).map_err(serde::de::Error::custom) + } +} + +/// Input supplied to an export-activation policy. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct ExportActivationRequest { + /// Kind of local or remote export target being considered. + pub target_kind: ExportActivationTargetKind, + /// Opaque target-local policy configuration. + #[serde(default)] + pub config: Json, +} + +/// Decision returned by an export-activation policy. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum ExportActivationDecision { + /// Construct and activate the exporter target. + Allow, + /// Suppress the exporter target for this activation. + Deny, +} + +/// Activation policy attached to one export target. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct ExportActivationPolicyConfig { + /// Plugin kind or dynamic manifest identifier that owns the callback. + pub provider: String, + /// Maximum policy evaluation time in milliseconds. + #[serde(default = "default_export_activation_timeout_millis")] + pub timeout_millis: u64, + /// Opaque target-local configuration passed to the policy. + #[serde(default)] + pub config: Json, +} + +/// Deferred export target registered by one plugin component. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct ExportTargetRegistration { + /// Component-local identifier, unique within one component instance. + pub id: String, + /// Namespaced kind presented to the policy provider. + pub target_kind: ExportActivationTargetKind, + /// Optional policy controlling target construction. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub activation_policy: Option, +} + +/// Minimum accepted policy timeout. +pub const MIN_EXPORT_ACTIVATION_TIMEOUT_MILLIS: u64 = 1_000; +/// Default policy timeout. +pub const DEFAULT_EXPORT_ACTIVATION_TIMEOUT_MILLIS: u64 = 30_000; +/// Maximum accepted policy timeout. +pub const MAX_EXPORT_ACTIVATION_TIMEOUT_MILLIS: u64 = 300_000; + +const fn default_export_activation_timeout_millis() -> u64 { + DEFAULT_EXPORT_ACTIVATION_TIMEOUT_MILLIS +} /// Diagnostic severity returned by plugin validation. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] diff --git a/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto b/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto index f2bf6e659..8a5471978 100644 --- a/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto +++ b/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto @@ -45,6 +45,8 @@ enum RegistrationSurface { REGISTRATION_SURFACE_UNSPECIFIED = 0; SUBSCRIBER = 1; EVENT_METADATA_INJECTOR = 2; + EXPORT_ACTIVATION_POLICY = 3; + EXPORT_TARGET = 4; TOOL_SANITIZE_REQUEST_GUARDRAIL = 10; TOOL_SANITIZE_RESPONSE_GUARDRAIL = 11; TOOL_CONDITIONAL_EXECUTION_GUARDRAIL = 12; @@ -150,6 +152,7 @@ message Registration { int32 priority = 3; bool break_chain = 4; reserved 5; + JsonEnvelope export_target = 6; } message InvokeRequest { @@ -165,6 +168,8 @@ message InvokeRequest { JsonEnvelope event = 10; ToolInvocation tool = 11; LlmInvocation llm = 12; + JsonEnvelope export_activation = 13; + JsonEnvelope export_target = 14; } } diff --git a/crates/worker/README.md b/crates/worker/README.md index 5b29a2daa..b710fd534 100644 --- a/crates/worker/README.md +++ b/crates/worker/README.md @@ -41,7 +41,8 @@ outcome fields now use structural protobuf messages. - **`WorkerPlugin`**: The plugin identity, validation, and registration contract. -- **`PluginContext`**: Typed registrations for all supported worker surfaces. +- **`PluginContext`**: Typed registrations for all supported worker surfaces, + including activation-time policies for plugin-managed exporters. - **`PluginRuntime` and continuations**: Host-runtime callbacks and tool/LLM execution-chain helpers. - **Canonical tool results**: `ToolNext` returns `ToolExecutionResult`, so diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 6c8d75742..cf4bb47f7 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -51,7 +51,10 @@ pub use nemo_relay_types::codec::optimization::{ }; pub use nemo_relay_types::codec::request::{ANNOTATED_LLM_REQUEST_SCHEMA, AnnotatedLlmRequest}; pub use nemo_relay_types::codec::response::AnnotatedLlmResponse; -pub use nemo_relay_types::plugin::{ConfigDiagnostic, DiagnosticLevel}; +pub use nemo_relay_types::plugin::{ + ConfigDiagnostic, DiagnosticLevel, ExportActivationDecision, ExportActivationPolicyConfig, + ExportActivationRequest, ExportActivationTargetKind, ExportTargetRegistration, +}; use nemo_relay_worker_proto::v1::plugin_worker_server::{PluginWorker, PluginWorkerServer}; use nemo_relay_worker_proto::v1::relay_host_runtime_client::RelayHostRuntimeClient; use nemo_relay_worker_proto::v1::{ @@ -123,6 +126,8 @@ pub type BoxFutureResult = Pin> + Send>>; pub type JsonStream = Pin> + Send>>; const JSON_SCHEMA: &str = "nemo.relay.Json@1"; +const EXPORT_ACTIVATION_REQUEST_SCHEMA: &str = "nemo.relay.ExportActivationRequest@1"; +const EXPORT_TARGET_REGISTRATION_SCHEMA: &str = "nemo.relay.ExportTargetRegistration@1"; const DATA_SCHEMA_SCHEMA: &str = "nemo.relay.DataSchema@1"; const LLM_REQUEST_SCHEMA: &str = "nemo.relay.LlmRequest@1"; @@ -184,6 +189,9 @@ type EventSanitizeFn = Arc BoxFutureResult + Send + Sync>; type EventMetadataInjectorFn = Arc) -> BoxFutureResult> + Send + Sync>; +type ExportActivationPolicyFn = + Arc BoxFutureResult + Send + Sync>; +type ExportTargetFn = Arc BoxFutureResult<()> + Send + Sync>; type ToolSanitizeFn = Arc BoxFutureResult + Send + Sync>; type ToolConditionalFn = Arc BoxFutureResult> + Send + Sync>; type ToolRequestFn = Arc BoxFutureResult + Send + Sync>; @@ -309,6 +317,8 @@ struct WorkerHandlers { registrations: Vec, subscribers: HashMap, event_metadata_injectors: HashMap, + export_activation_policies: HashMap, + export_targets: HashMap, mark_sanitizers: HashMap, scope_start_sanitizers: HashMap, scope_end_sanitizers: HashMap, @@ -386,6 +396,45 @@ impl PluginContext { ); } + /// Registers this plugin's activation-time policy for exporters. + pub fn register_export_activation_policy(&mut self, callback: F) + where + F: Fn(ExportActivationRequest) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + const NAME: &str = "export_activation_policy"; + self.push_registration(NAME, RegistrationSurface::ExportActivationPolicy, 0, false); + self.handlers.export_activation_policies.insert( + NAME.into(), + Arc::new(move |request| Box::pin(callback(request))), + ); + } + + /// Registers a deferred local or remote export target. + pub fn register_export_target( + &mut self, + registration: ExportTargetRegistration, + activate: F, + ) where + F: Fn() -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let name = registration.id.clone(); + self.handlers.registrations.push(Registration { + local_name: name.clone(), + surface: RegistrationSurface::ExportTarget as i32, + priority: 0, + break_chain: false, + export_target: Some(infallible_json_envelope( + EXPORT_TARGET_REGISTRATION_SCHEMA, + ®istration, + )), + }); + self.handlers + .export_targets + .insert(name, Arc::new(move || Box::pin(activate()))); + } + fn register_event_sanitizer( &mut self, name: &str, @@ -741,6 +790,7 @@ impl PluginContext { surface: surface as i32, priority, break_chain, + export_target: None, }); } } @@ -1819,6 +1869,13 @@ impl WorkerService { self.invoke_event_metadata_injector_response(request, &scope) .await } + RegistrationSurface::ExportActivationPolicy => { + self.invoke_export_activation_policy_response(request, &scope) + .await + } + RegistrationSurface::ExportTarget => { + self.invoke_export_target_response(request, &scope).await + } RegistrationSurface::MarkSanitizeGuardrail | RegistrationSurface::ScopeSanitizeStartGuardrail | RegistrationSurface::ScopeSanitizeEndGuardrail => { @@ -1871,6 +1928,44 @@ impl WorkerService { )?)) } + async fn invoke_export_activation_policy_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = match request.payload { + Some(nemo_relay_worker_proto::v1::invoke_request::Payload::ExportActivation(value)) => { + value + } + _ => { + return Err(WorkerSdkError::InvalidInput( + "export activation policy requires an export_activation payload".into(), + )); + } + }; + if payload.schema != EXPORT_ACTIVATION_REQUEST_SCHEMA { + return Err(WorkerSdkError::InvalidInput(format!( + "unsupported export activation request schema {:?}", + payload.schema + ))); + } + let policy_request = decode_json_envelope::(&payload)?; + let handler = self.export_activation_policy(&request.registration_name)?; + let future = with_thread_scope(scope, || handler(policy_request)); + Ok(json_response(serde_json::to_value(future.await?)?)) + } + + async fn invoke_export_target_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let handler = self.export_target(&request.registration_name)?; + let future = with_thread_scope(scope, || handler()); + future.await?; + Ok(json_response(Json::Null)) + } + async fn invoke_event_sanitize_response( &self, request: InvokeRequest, @@ -2116,6 +2211,32 @@ impl WorkerService { }) } + fn export_activation_policy(&self, name: &str) -> Result { + self.handlers + .lock() + .map_err(|err| WorkerSdkError::Callback(format!("handler lock poisoned: {err}")))? + .export_activation_policies + .get(name) + .cloned() + .ok_or_else(|| { + WorkerSdkError::InvalidInput(format!( + "export activation policy '{name}' not registered" + )) + }) + } + + fn export_target(&self, name: &str) -> Result { + self.handlers + .lock() + .map_err(|err| WorkerSdkError::Callback(format!("handler lock poisoned: {err}")))? + .export_targets + .get(name) + .cloned() + .ok_or_else(|| { + WorkerSdkError::InvalidInput(format!("export target '{name}' not registered")) + }) + } + fn event_sanitizer(&self, surface: RegistrationSurface, name: &str) -> Result { let handlers = self .handlers @@ -2743,6 +2864,8 @@ fn all_surfaces() -> Vec { vec![ RegistrationSurface::Subscriber, RegistrationSurface::EventMetadataInjector, + RegistrationSurface::ExportActivationPolicy, + RegistrationSurface::ExportTarget, RegistrationSurface::ToolSanitizeRequestGuardrail, RegistrationSurface::ToolSanitizeResponseGuardrail, RegistrationSurface::ToolConditionalExecutionGuardrail, diff --git a/crates/worker/tests/worker_sdk_tests.rs b/crates/worker/tests/worker_sdk_tests.rs index d2e0717d9..0c6361f5d 100644 --- a/crates/worker/tests/worker_sdk_tests.rs +++ b/crates/worker/tests/worker_sdk_tests.rs @@ -20,11 +20,12 @@ use futures_util::{Stream, StreamExt}; use hyper_util::rt::TokioIo; use nemo_relay_types::api::event::{BaseEvent, Event, MarkEvent, PendingMarkSpec}; use nemo_relay_worker::{ - ANNOTATED_LLM_REQUEST_SCHEMA, DataSchema, EmitMarkOptions, Json, JsonStream, LlmNext, - LlmRequest, LlmStreamNext, LogSeverity, MetricKind, MetricMeasurement, MetricValueType, - PluginContext, PluginRuntime, Result, ScopeType, ToolExecutionInterceptOutcome, ToolNext, - WorkerPlugin, WorkerSdkError, WorkerServerConfig, serve_plugin, serve_plugin_arc, - serve_plugin_arc_with_config, + ANNOTATED_LLM_REQUEST_SCHEMA, DataSchema, EmitMarkOptions, ExportActivationDecision, + ExportActivationPolicyConfig, ExportActivationTargetKind, ExportTargetRegistration, Json, + JsonStream, LlmNext, LlmRequest, LlmStreamNext, LogSeverity, MetricKind, MetricMeasurement, + MetricValueType, PluginContext, PluginRuntime, Result, ScopeType, + ToolExecutionInterceptOutcome, ToolNext, WorkerPlugin, WorkerSdkError, WorkerServerConfig, + serve_plugin, serve_plugin_arc, serve_plugin_arc_with_config, }; use nemo_relay_worker_proto::v1::plugin_worker_client::PluginWorkerClient; use nemo_relay_worker_proto::v1::relay_host_runtime_server::{ @@ -127,6 +128,16 @@ async fn worker_service_enforces_auth_and_reports_registrations() { .supported_surfaces .contains(&(RegistrationSurface::EventMetadataInjector as i32)) ); + assert!( + handshake + .supported_surfaces + .contains(&(RegistrationSurface::ExportActivationPolicy as i32)) + ); + assert!( + handshake + .supported_surfaces + .contains(&(RegistrationSurface::ExportTarget as i32)) + ); let bad_health = client .health(Request::new(HealthRequest { @@ -213,7 +224,7 @@ async fn worker_service_enforces_auth_and_reports_registrations() { assert_eq!(invalid_register_config.code(), tonic::Code::InvalidArgument); let registrations = register_plugin(&mut client).await; - assert_eq!(registrations.len(), 22); + assert_eq!(registrations.len(), 24); for local_name in [ "llm-sanitize-request", "llm-sanitize-response", @@ -545,6 +556,11 @@ async fn worker_service_invokes_every_registration_surface() { .await; assert_eq!(metadata, json!({"worker.event_name": "subscriber-event"})); + let policy = invoke_json(&mut client, export_activation_invoke()).await; + assert_eq!(policy, json!("allow")); + let target = invoke_json(&mut client, export_target_invoke()).await; + assert_eq!(target, Json::Null); + let mark_fields = invoke_json( &mut client, event_invoke_surface("event-sanitize", RegistrationSurface::MarkSanitizeGuardrail), @@ -1833,6 +1849,29 @@ impl WorkerPlugin for SurfacePlugin { json!(event.name()), )])) }); + ctx.register_export_activation_policy(|request| async move { + Ok( + if request.target_kind == nemo_relay_worker::ExportActivationTargetKind::OtlpTrace + && request.config == json!({"enabled": true}) + { + ExportActivationDecision::Allow + } else { + ExportActivationDecision::Deny + }, + ) + }); + ctx.register_export_target( + ExportTargetRegistration { + id: "self-otel".into(), + target_kind: ExportActivationTargetKind::OTLP_TRACE, + activation_policy: Some(ExportActivationPolicyConfig { + provider: PLUGIN_ID.into(), + timeout_millis: 30_000, + config: json!({"enabled": true}), + }), + }, + || async { Ok(()) }, + ); ctx.register_mark_sanitize_guardrail("event-sanitize", 1, |event, mut fields| { let event_name = event.name().to_owned(); async move { @@ -2688,6 +2727,45 @@ fn event_invoke_surface(registration_name: &str, surface: RegistrationSurface) - } } +fn export_activation_invoke() -> InvokeRequest { + InvokeRequest { + activation_id: ACTIVATION_ID.into(), + invocation_id: "invoke-policy".into(), + registration_name: "export_activation_policy".into(), + surface: RegistrationSurface::ExportActivationPolicy as i32, + continuation_id: String::new(), + scope: None, + auth_token: AUTH_TOKEN.into(), + payload: Some( + nemo_relay_worker_proto::v1::invoke_request::Payload::ExportActivation( + json_envelope( + "nemo.relay.ExportActivationRequest@1", + &json!({"target_kind": "nemo_relay.otlp.trace", "config": {"enabled": true}}), + ) + .expect("encode export activation request"), + ), + ), + } +} + +fn export_target_invoke() -> InvokeRequest { + InvokeRequest { + activation_id: ACTIVATION_ID.into(), + invocation_id: "invoke-target".into(), + registration_name: "self-otel".into(), + surface: RegistrationSurface::ExportTarget as i32, + continuation_id: String::new(), + scope: None, + auth_token: AUTH_TOKEN.into(), + payload: Some( + nemo_relay_worker_proto::v1::invoke_request::Payload::ExportTarget( + json_envelope("nemo.relay.ExportTargetActivation@1", &Json::Null) + .expect("encode export target activation"), + ), + ), + } +} + fn tool_invoke( registration_name: &str, surface: RegistrationSurface, diff --git a/docs/build-plugins/dynamic-plugins/about.mdx b/docs/build-plugins/dynamic-plugins/about.mdx index 5a2847f5a..a3ee82461 100644 --- a/docs/build-plugins/dynamic-plugins/about.mdx +++ b/docs/build-plugins/dynamic-plugins/about.mdx @@ -70,6 +70,14 @@ operators must enable a registered dynamic plugin explicitly. Declare only the capabilities the plugin needs. Add `config_schema.path` only with the `config_schema` capability. +An export-activation provider also declares `export_activation_policy` in +`capabilities.items`. Relay recognizes that capability before activating +static observability components, starts the provider early, and evaluates all +targets that reference the provider manifest ID. A provider registers exactly +one callback; duplicate callbacks under the same manifest ID are rejected. +Other middleware and subscribers registered by that component are installed +during the same early registration and retain their configured priorities. + The following requirements vary by execution lane: | Manifest area | Native dynamic plugin | Worker plugin | diff --git a/docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx b/docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx index bcb836f96..52c44e947 100644 --- a/docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx +++ b/docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx @@ -211,6 +211,57 @@ returns at most 32 entries. Entries do not identify the emitting plugin. If a worker connects to an older Relay host, the SDK translates gRPC `UNIMPLEMENTED` into an unsupported-runtime-diagnostics error. +## Register an Export Activation Policy + +Add `"export_activation_policy"` to the worker manifest +`capabilities.items`. Worker protocol `grpc-v1` carries the versioned request +and decision through its existing unary `Invoke` RPC: + +```python +from nemo_relay_plugin import ( + ExportActivationDecision, + ExportActivationPolicyConfig, + ExportActivationRequest, + ExportActivationTargetKind, + ExportTargetRegistration, +) + + +async def decide_export(request: ExportActivationRequest) -> ExportActivationDecision: + config = request.config if isinstance(request.config, dict) else {} + country = await lookup_current_public_ip_country() + allowed_countries = config.get("allowed_countries", []) + return ( + ExportActivationDecision.ALLOW + if country in allowed_countries + else ExportActivationDecision.DENY + ) + + +ctx.register_export_activation_policy(decide_export) +ctx.register_export_target( + ExportTargetRegistration( + id="audit-otlp", + target_kind=ExportActivationTargetKind.OTLP_TRACE, + activation_policy=ExportActivationPolicyConfig( + provider="examples.python_grpc_worker", + timeout_millis=30_000, + config={"allowed_countries": ["US", "CA"]}, + ), + ), + start_custom_otlp_exporter, +) +``` + +Relay starts a referenced policy worker before constructing gated exporters. +It sends only the target kind and opaque policy configuration, never exporter +destinations, headers, or credentials. If the worker does not advertise the +surface, exits, errors, or exceeds the per-target timeout, Relay denies that +target without preventing the rest of the runtime from starting. The provider +above is the same worker manifest that owns `start_custom_otlp_exporter`, so the +worker starts as bootstrap infrastructure while the OTLP exporter does not +start unless its policy allows it. + ## Emit Typed Marks and Metrics Worker plugins can assign a schema and OTLP log severity to a mark without diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example.mdx index 912b65cb6..3ce48e582 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example.mdx @@ -131,3 +131,55 @@ Keep tests outside `src`, use the shared DTOs re-exported by `nemo-relay-plugin`, and do not retain host-owned ABI handles after a callback returns. Refer to [Configure Discoverable Plugins](/configure-plugins/discoverable-plugins) for trust-policy configuration. + +## Register an Export Activation Policy + +Native ABI v4 exposes the activation policy through the existing asynchronous +registration table, so it adds no fields to the ABI-v4 descriptor. Add +`"export_activation_policy"` to the manifest `capabilities.items`, keep the +Relay compatibility range at `>=0.8.0,<1.0`, and register one callback: + +```rust +use nemo_relay_plugin::{ + ExportActivationDecision, ExportActivationPolicyConfig, ExportActivationTargetKind, + ExportTargetRegistration, PluginContext, Result, +}; + +fn register_geo_gated_otlp(context: &mut PluginContext<'_>) -> Result<()> { + context.register_export_activation_policy(|request| async move { + let country = lookup_current_public_ip_country().await?; + let allowed = request.config["allowed_countries"] + .as_array() + .is_some_and(|countries| { + countries + .iter() + .any(|value| value.as_str() == Some(country.as_str())) + }); + Ok(if allowed { + ExportActivationDecision::Allow + } else { + ExportActivationDecision::Deny + }) + })?; + + context.register_export_target( + ExportTargetRegistration { + id: "audit-otlp".into(), + target_kind: ExportActivationTargetKind::OTLP_TRACE, + activation_policy: Some(ExportActivationPolicyConfig { + provider: "examples.rust_native_policy".into(), + timeout_millis: 30_000, + config: serde_json::json!({"allowed_countries": ["US", "CA"]}), + }), + }, + || async { start_custom_otlp_exporter().await }, + ) +} +``` + +The callback can implement GeoIP, available-memory, consent, or environment +checks. Keep it bounded and deterministic for the activation attempt. Relay +applies the target's timeout and fails that target closed if the callback +returns an error. In this example, `examples.rust_native_policy` is both the dynamic +plugin's manifest ID and the provider referenced by its own deferred OTLP +export target. Construct the exporter only in `start_custom_otlp_exporter`. diff --git a/docs/build-plugins/language-binding/register-behavior.mdx b/docs/build-plugins/language-binding/register-behavior.mdx index 205d03a86..0ca0e0a52 100644 --- a/docs/build-plugins/language-binding/register-behavior.mdx +++ b/docs/build-plugins/language-binding/register-behavior.mdx @@ -42,6 +42,57 @@ Use the context methods so component name qualification and rollback also apply to these registries. Refer to [Event Sanitizers](/reference/event-sanitizers) for the binding-specific method names. +## Gate Plugin Exporters During Activation + +Export activation hooks are component-scoped lifecycle APIs, not middleware. +They do not have priorities, scope inheritance, or a place in tool, LLM, or +event execution chains. + +A plugin can register one policy provider and any number of deferred export +targets. Register both from the same component when an exporter should use its +own policy. Relay completes component registration first, evaluates each policy +once, and invokes only allowed or ungated target callbacks. Construct files, +connections, exporter workers, and background tasks inside the target callback. +Relay retains an allowed callback until activation teardown so callback-owned +state remains alive. + +Policy requests contain only a validated namespaced `target_kind` and the +target's opaque `config`. Relay does not include target IDs, paths, URLs, +headers, or credentials. Missing providers, callback errors, and timeouts deny +only the affected target. The timeout defaults to 30 seconds and accepts values +from 1 second through 5 minutes. + +```python +async def policy(request): + country = await lookup_current_public_ip_country() + return "allow" if country in request["config"]["allowed_countries"] else "deny" + +async def start_exporter(): + # Keep the collector URL and credentials in component config, not policy config. + exporter = create_otlp_trace_exporter(component_config["otlp"]) + await exporter.start() + +context.register_export_activation_policy(policy) +context.register_export_target( + { + "id": "audit-otlp", + "target_kind": "nemo_relay.otlp.trace", + "activation_policy": { + "provider": "com.example.audit", + "timeout_millis": 30_000, + "config": {"allowed_countries": ["US", "CA"]}, + }, + }, + start_exporter, +) +``` + +Dynamic policy providers must declare the `export_activation_policy` manifest +capability. Worker plugins must also advertise both activation-hook surfaces +that they register. The provider library or worker is bootstrap infrastructure: +Relay loads or starts it before policy evaluation, but it does not construct a +gated exporter before the policy allows it. + ## Header Plugin Example The same model applies in every binding: validate component-local config, then install middleware through the component-scoped registration context. diff --git a/docs/configure-plugins/observability/atif.mdx b/docs/configure-plugins/observability/atif.mdx index bb1602879..cbe6584fb 100644 --- a/docs/configure-plugins/observability/atif.mdx +++ b/docs/configure-plugins/observability/atif.mdx @@ -69,6 +69,7 @@ The following table describes the top-level ATIF settings: | `extra` | Omitted | Optional ATIF agent metadata. | | `output_directory` | Current working directory | Directory containing trajectory files and the recovery copy when every configured remote destination fails. | | `filename_template` | `nemo-relay-atif-{session_id}.json` | Must contain `{session_id}`. Can contain `{metadata.}` placeholders for metadata-based routing. | +| `local_activation_policy` | Omitted | Optional policy for the implicit local file and local recovery fallback. A denial suppresses both uses. | | `storage` | Omitted | Optional list of remote storage destinations. When non-empty, trajectories are uploaded to every configured backend. If all fail for one trajectory, Relay writes a local recovery copy. Refer to [Remote Storage](#remote-storage). | ### Metadata-Based Paths @@ -120,6 +121,12 @@ trajectory to every configured backend instead of writing a local file. for a local recovery copy when every configured remote destination fails for a trajectory. +Each HTTP or S3 entry can include `activation_policy`. Relay evaluates entries +independently and preserves their original configuration indices in +diagnostics. If every configured remote entry is denied, Relay disables the +dispatcher instead of promoting the implicit local destination to primary +storage. + Each storage entry is tagged with a `type` discriminator so additional backends can be added without breaking existing configs. S3-compatible object storage and HTTP endpoints are supported. diff --git a/docs/configure-plugins/observability/atof.mdx b/docs/configure-plugins/observability/atof.mdx index 1921047cb..74f6eda01 100644 --- a/docs/configure-plugins/observability/atof.mdx +++ b/docs/configure-plugins/observability/atof.mdx @@ -68,6 +68,11 @@ The following table describes the top-level ATOF settings: | `enabled` | `false` | Must be `true` to export events. | | `sinks` | `[]` | File and stream destinations. An enabled ATOF section requires at least one sink. Each event delivered to the ATOF subscriber is sent to every configured sink. | +Each file or stream sink can include an `activation_policy` table. Denying a +file sink prevents Relay from opening its file. If every sink is denied, Relay +does not install an ATOF exporter; canonical event generation and unrelated +subscribers continue normally. + ## File Sinks Each file sink writes the raw event stream to one JSONL file. The following diff --git a/docs/configure-plugins/observability/configuration.mdx b/docs/configure-plugins/observability/configuration.mdx index 3face373c..04d0e76d7 100644 --- a/docs/configure-plugins/observability/configuration.mdx +++ b/docs/configure-plugins/observability/configuration.mdx @@ -89,6 +89,54 @@ Destination collision validation is signal-aware. Duplicate destinations within logs or within metrics are rejected, while different OTLP signals can share an authority. +## Export Activation Policies + +Version 4 export targets accept an optional `activation_policy`. The +`provider` is the ordinary plugin kind or dynamic-plugin manifest ID of a +component that registers one policy callback. Dynamic providers also declare +the `export_activation_policy` capability: + +```toml +[components.config.opentelemetry.endpoints.activation_policy] +provider = "com.example.runtime-policy" +timeout_millis = 30000 +config = { allowed_countries = ["US", "CA"], min_available_memory_mib = 2048 } +``` + +Relay evaluates the policy once for each configured target before constructing +that exporter. The callback receives only the target kind and the opaque +`config` value. Relay does not send the destination URL, headers, credentials, +or other exporter settings to the policy. The decision remains fixed until +plugin reconfiguration or restart. + +Configuration layering finishes before policy providers register. Each +resolved plugin activation owns one provider namespace, and every target in +that activation resolves its provider manifest ID against that namespace. +Provider implementations do not shadow one another by configuration layer; +layered targets can reference the same provider with different opaque `config` +values. A later activation receives a new provider namespace. + +Policies are available on OTLP trace, log, and metric endpoints; ATOF stream +sinks; and ATIF HTTP and S3 targets. Local ATOF and ATIF file destinations do +not accept policies. Derived OTLP log and metric endpoints inherit the trace +policy, but Relay evaluates each signal separately with `otlp_log` or +`otlp_metric` as its target kind. + +The timeout defaults to 30000 milliseconds and must be between 1000 and 300000 +(one second and five minutes). Relay enforces these bounds when evaluating a +policy even when the configuration policy reports unsupported values as +warnings. +A deny decision, missing provider, callback error, unsupported worker surface, +or timeout denies only that target. Relay continues activation and emits a +structured diagnostic that contains the provider ID, target kind, and reason, +but no destination or credential data. A denied target is never constructed, +connected, or given a background exporter task. + +Mixed ATOF configurations retain allowed local sinks. When policy evaluation +denies every remote ATIF target, Relay disables the ATIF dispatcher and does +not silently fall back to its legacy local-file destination. Policies are an +activation-time control only; per-event reevaluation is not supported. + ## Full LLM Payloads By default, repeated LLM start events contain only the current user turn. Set @@ -163,6 +211,7 @@ from nemo_relay.observability import ( ObservabilityConfig, OpenTelemetryLogSectionConfig, OpenTelemetryMetricSectionConfig, + ExportActivationPolicyConfig, OpenTelemetrySectionConfig, OpenTelemetryEndpointConfig, ) @@ -178,6 +227,10 @@ config = plugin.PluginConfig( type="gen_ai", endpoint="http://localhost:4318/v1/traces", service_name="agent-service", + activation_policy=ExportActivationPolicyConfig( + provider="com.example.runtime-policy", + config={"allowed_countries": ["US", "CA"]}, + ), ) ], logs=OpenTelemetryLogSectionConfig(enabled=True), @@ -203,6 +256,10 @@ const component = observability.ComponentSpec({ type: "gen_ai", endpoint: "http://localhost:4318/v1/traces", service_name: "agent-service", + activation_policy: observability.exportActivationPolicy({ + provider: "com.example.runtime-policy", + config: { allowed_countries: ["US", "CA"] }, + }), }), ], logs: observability.openTelemetryLogConfig({ enabled: true }), diff --git a/docs/configure-plugins/observability/opentelemetry.mdx b/docs/configure-plugins/observability/opentelemetry.mdx index 8ea56d230..f3cda133a 100644 --- a/docs/configure-plugins/observability/opentelemetry.mdx +++ b/docs/configure-plugins/observability/opentelemetry.mdx @@ -95,6 +95,11 @@ max_queue_size = 4096 max_export_batch_size = 512 scheduled_delay_millis = 1000 +[components.config.opentelemetry.endpoints.activation_policy] +provider = "com.example.runtime-policy" +timeout_millis = 30000 +config = { allowed_countries = ["US", "CA"] } + [components.config.opentelemetry.endpoints.header_env] authorization = "OTEL_AUTHORIZATION" @@ -144,6 +149,7 @@ to the other exporters. | `mark_exclude_names` | `["llm.chunk"]` | Mark names excluded from `full` and `openinference` projection. | | `attribute_mappings` | `[]` | `{ key, alias }` copies applied by `full` and `openinference` projection. | | `promote_metadata_prefixes` | `[]` | Literal prefixes that select sanitized Event metadata to copy to top-level span attributes. | +| `activation_policy` | Omitted | Optional activation hook for this endpoint. Relay evaluates derived trace, log, and metric targets separately and does not construct denied exporters. | Metadata promotion preserves the original Event metadata and does not replace attributes produced by a trace projection or `attribute_mappings`. Relay omits diff --git a/examples/python-grpc-worker-plugin/README.md b/examples/python-grpc-worker-plugin/README.md index b24a3cb6c..75929a5d0 100644 --- a/examples/python-grpc-worker-plugin/README.md +++ b/examples/python-grpc-worker-plugin/README.md @@ -8,7 +8,9 @@ SPDX-License-Identifier: Apache-2.0 This example shows a Python worker plugin using the `nemo-relay-plugin` SDK. It registers tool request and execution intercepts, calls the host continuation, preserves the upstream tool-result annotation, and emits marks through both the -host runtime and the execution outcome. +host runtime and the execution outcome. It also registers an export-activation +policy that allows an export target only when its opaque policy config contains +`allow = true`. The example targets Relay 0.8 or later and the canonical `grpc-v1` tool-result contract declared in `relay-plugin.toml`. diff --git a/examples/python-grpc-worker-plugin/nemo_relay_python_grpc_worker_example/worker.py b/examples/python-grpc-worker-plugin/nemo_relay_python_grpc_worker_example/worker.py index 4fe25c728..d5a7f740c 100644 --- a/examples/python-grpc-worker-plugin/nemo_relay_python_grpc_worker_example/worker.py +++ b/examples/python-grpc-worker-plugin/nemo_relay_python_grpc_worker_example/worker.py @@ -10,6 +10,8 @@ from nemo_relay_plugin import ( ConfigDiagnostic, DiagnosticLevel, + ExportActivationDecision, + ExportActivationRequest, Json, PendingMarkSpec, PluginContext, @@ -66,6 +68,16 @@ def register(self, ctx: PluginContext, config: Json) -> None: if not isinstance(tag, str): raise TypeError("tag must be a string") + async def decide_export( + request: ExportActivationRequest, + ) -> ExportActivationDecision: + # Each endpoint opts in with `config = { allow = true }`. A real + # provider can replace this with GeoIP, memory, consent, or other + # activation-time logic without receiving the endpoint URL or secrets. + if isinstance(request.config, dict) and request.config.get("allow") is True: + return ExportActivationDecision.ALLOW + return ExportActivationDecision.DENY + async def inject_event_metadata(event: Any) -> dict[str, Json]: if config.get("event_metadata_injector_error") is True: raise RuntimeError("Python Event metadata injector error requested") @@ -101,6 +113,7 @@ async def tag_tool_execution( ], ) + ctx.register_export_activation_policy(decide_export) ctx.register_event_metadata_injector("example_event_metadata_injector", inject_event_metadata) if config.get("event_metadata_injector_only") is True: return diff --git a/examples/python-grpc-worker-plugin/relay-plugin.toml b/examples/python-grpc-worker-plugin/relay-plugin.toml index 804ac8acf..8b6d38555 100644 --- a/examples/python-grpc-worker-plugin/relay-plugin.toml +++ b/examples/python-grpc-worker-plugin/relay-plugin.toml @@ -15,14 +15,14 @@ worker_protocol = "grpc-v1" enabled = false [capabilities] -items = ["plugin_worker"] +items = ["plugin_worker", "export_activation_policy"] [source] manifest_root = "." artifact = "nemo_relay_python_grpc_worker_example/worker.py" [integrity] -sha256 = "sha256:7bcc2fb993336c84fc319b0c5669c83ee9257822160ab258e2f15d2429fbd97b" +sha256 = "sha256:46b0af6d07932cedf1737a11750a2cfe79b5c7d155237a8ddc36b845620cdc88" [load] runtime = "python" diff --git a/examples/rust-native-plugin/README.md b/examples/rust-native-plugin/README.md index 790605ffa..65fa9b039 100644 --- a/examples/rust-native-plugin/README.md +++ b/examples/rust-native-plugin/README.md @@ -100,6 +100,8 @@ nemo-relay --bind 127.0.0.1:4040 The example registers the following runtime behavior: +- An export-activation policy that allows an export target only when its opaque + policy config contains `allow = true`. - A subscriber that emits a mark when it sees non-plugin scope starts. - Tool sanitize request/response guardrails for observability payload tagging. - Conditional execution guardrails for tools and LLMs controlled by config. diff --git a/examples/rust-native-plugin/relay-plugin.toml b/examples/rust-native-plugin/relay-plugin.toml index 8584ad638..2ea311d21 100644 --- a/examples/rust-native-plugin/relay-plugin.toml +++ b/examples/rust-native-plugin/relay-plugin.toml @@ -15,7 +15,7 @@ native_api = "1" enabled = false [capabilities] -items = ["plugin_native", "config_schema"] +items = ["plugin_native", "config_schema", "export_activation_policy"] [config_schema] path = "config.schema.json" diff --git a/examples/rust-native-plugin/src/lib.rs b/examples/rust-native-plugin/src/lib.rs index 307c01729..062fdf7df 100644 --- a/examples/rust-native-plugin/src/lib.rs +++ b/examples/rust-native-plugin/src/lib.rs @@ -4,9 +4,9 @@ use futures::StreamExt; use nemo_relay_plugin::{ CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, Json, - LlmJsonAsyncStream, LlmRequest, LlmRequestInterceptOutcome, NativeExecutorConfig, NativePlugin, - PendingMarkSpec, PluginContext, PluginRuntime, ScopeCategory, ScopeType, - ToolExecutionInterceptOutcome, + ExportActivationDecision, LlmJsonAsyncStream, LlmRequest, LlmRequestInterceptOutcome, + NativeExecutorConfig, NativePlugin, PendingMarkSpec, PluginContext, PluginRuntime, ScopeCategory, + ScopeType, ToolExecutionInterceptOutcome, }; use serde_json::{Map, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -186,6 +186,18 @@ impl NativePlugin for ExampleNativePlugin { let config = ExampleConfig::parse(plugin_config)?; let runtime = ctx.runtime(); + ctx.register_export_activation_policy(|request| async move { + // This example is intentionally generic: each endpoint opts in by + // passing `config = { allow = true }` in its activation policy. + // A production provider could make the same decision from GeoIP, + // available memory, user consent, or another runtime signal. + Ok(if request.config.get("allow").and_then(Json::as_bool) == Some(true) { + ExportActivationDecision::Allow + } else { + ExportActivationDecision::Deny + }) + })?; + ctx.register_subscriber("example_native_subscriber", { let runtime = runtime.clone(); let tag = config.tag.clone(); diff --git a/go/nemo_relay/callbacks.go b/go/nemo_relay/callbacks.go index 4cdb8690a..a09565894 100644 --- a/go/nemo_relay/callbacks.go +++ b/go/nemo_relay/callbacks.go @@ -56,6 +56,8 @@ typedef char* (*NemoRelayToolExecNextFn)(const char* args_json, void* next_ctx); typedef char* (*NemoRelayToolExecInterceptCb)(void* user_data, const char* args_json, NemoRelayToolExecNextFn next_fn, void* next_ctx); typedef char* (*NemoRelayLlmExecNextFn)(const char* native_json, void* next_ctx); typedef char* (*NemoRelayLlmExecInterceptCb)(void* user_data, const char* native_json, NemoRelayLlmExecNextFn next_fn, void* next_ctx); +typedef char* (*NemoRelayExportActivationPolicyCb)(void* user_data, const char* request_json); +typedef int32_t (*NemoRelayExportTargetActivationCb)(void* user_data); // Helper to call the tool exec next function pointer from Go static inline char* callToolExecNext(NemoRelayToolExecNextFn next_fn, const char* args_json, void* next_ctx) { @@ -684,6 +686,37 @@ func goFreeTrampoline(userData unsafe.Pointer) { unregisterClosure(userData) } +//export goExportActivationPolicyTrampoline +func goExportActivationPolicyTrampoline(userData unsafe.Pointer, requestJSON *C.char) *C.char { + fn := lookupClosure(userData).(ExportActivationPolicyFunc) + var request ExportActivationRequest + if err := jsonUnmarshal([]byte(C.GoString(requestJSON)), &request); err != nil { + setLastErrorMessage(err.Error()) + return nil + } + decision, err := fn(request) + if err != nil { + setLastErrorMessage(err.Error()) + return nil + } + payload, err := jsonMarshal(decision) + if err != nil { + setLastErrorMessage(err.Error()) + return nil + } + return C.CString(string(payload)) +} + +//export goExportTargetActivationTrampoline +func goExportTargetActivationTrampoline(userData unsafe.Pointer) C.int32_t { + fn := lookupClosure(userData).(ExportTargetActivationFunc) + if err := fn(); err != nil { + setLastErrorMessage(err.Error()) + return 5 // NemoRelayStatus::Internal + } + return 0 // NemoRelayStatus::Ok +} + //export goLlmRequestTrampoline func goLlmRequestTrampoline(userData unsafe.Pointer, request *C.FfiLLMRequest, context C.NemoRelayLlmSanitizeRequestContext) *C.FfiLLMRequest { fn := lookupClosure(userData).(LLMRequestFunc) diff --git a/go/nemo_relay/observability_plugin.go b/go/nemo_relay/observability_plugin.go index 987d71cfd..2a540ce3b 100644 --- a/go/nemo_relay/observability_plugin.go +++ b/go/nemo_relay/observability_plugin.go @@ -8,6 +8,13 @@ import "encoding/json" // ObservabilityPluginKind is the top-level plugin kind used by the core observability component. const ObservabilityPluginKind = "observability" +// ExportActivationPolicyConfig attaches one policy provider to an export target. +type ExportActivationPolicyConfig struct { + Provider string `json:"provider"` + TimeoutMillis uint64 `json:"timeout_millis,omitempty"` + Config any `json:"config,omitempty"` +} + // ObservabilityConfig is the canonical Go shape for the observability plugin config document. type ObservabilityConfig struct { Version uint32 `json:"version,omitempty"` @@ -29,16 +36,17 @@ type ObservabilityOpenTelemetryConfig struct { // ObservabilityOpenTelemetrySignalEndpointConfig configures one log or metric OTLP destination. type ObservabilityOpenTelemetrySignalEndpointConfig struct { - Endpoint string `json:"endpoint"` - Transport string `json:"transport,omitempty"` - Headers map[string]string `json:"headers,omitempty"` - HeaderEnv map[string]string `json:"header_env,omitempty"` - ResourceAttributes map[string]string `json:"resource_attributes,omitempty"` - ServiceName string `json:"service_name,omitempty"` - ServiceNamespace string `json:"service_namespace,omitempty"` - ServiceVersion string `json:"service_version,omitempty"` - InstrumentationScope string `json:"instrumentation_scope,omitempty"` - TimeoutMillis uint64 `json:"timeout_millis,omitempty"` + Endpoint string `json:"endpoint"` + Transport string `json:"transport,omitempty"` + Headers map[string]string `json:"headers,omitempty"` + HeaderEnv map[string]string `json:"header_env,omitempty"` + ResourceAttributes map[string]string `json:"resource_attributes,omitempty"` + ServiceName string `json:"service_name,omitempty"` + ServiceNamespace string `json:"service_namespace,omitempty"` + ServiceVersion string `json:"service_version,omitempty"` + InstrumentationScope string `json:"instrumentation_scope,omitempty"` + TimeoutMillis uint64 `json:"timeout_millis,omitempty"` + ActivationPolicy *ExportActivationPolicyConfig `json:"activation_policy,omitempty"` } // ObservabilityOpenTelemetryLogConfig configures the plugin's OTLP log pipeline. @@ -65,24 +73,25 @@ type ObservabilityOpenTelemetryMetricConfig struct { // ObservabilityOpenTelemetryEndpointConfig configures one typed OTLP destination. type ObservabilityOpenTelemetryEndpointConfig struct { - Type OpenTelemetryType `json:"type"` - Endpoint string `json:"endpoint"` - MarkProjection string `json:"mark_projection,omitempty"` - MarkExcludeNames []string `json:"mark_exclude_names,omitempty"` - AttributeMappings []OtlpAttributeMapping `json:"attribute_mappings,omitempty"` - PromoteMetadataPrefixes []string `json:"promote_metadata_prefixes,omitempty"` - Transport string `json:"transport,omitempty"` - Headers map[string]string `json:"headers,omitempty"` - HeaderEnv map[string]string `json:"header_env,omitempty"` - ResourceAttributes map[string]string `json:"resource_attributes,omitempty"` - ServiceName string `json:"service_name,omitempty"` - ServiceNamespace string `json:"service_namespace,omitempty"` - ServiceVersion string `json:"service_version,omitempty"` - InstrumentationScope string `json:"instrumentation_scope,omitempty"` - TimeoutMillis uint64 `json:"timeout_millis,omitempty"` - MaxQueueSize *uint64 `json:"max_queue_size,omitempty"` - MaxExportBatchSize *uint64 `json:"max_export_batch_size,omitempty"` - ScheduledDelayMillis *uint64 `json:"scheduled_delay_millis,omitempty"` + Type OpenTelemetryType `json:"type"` + Endpoint string `json:"endpoint"` + MarkProjection string `json:"mark_projection,omitempty"` + MarkExcludeNames []string `json:"mark_exclude_names,omitempty"` + AttributeMappings []OtlpAttributeMapping `json:"attribute_mappings,omitempty"` + PromoteMetadataPrefixes []string `json:"promote_metadata_prefixes,omitempty"` + Transport string `json:"transport,omitempty"` + Headers map[string]string `json:"headers,omitempty"` + HeaderEnv map[string]string `json:"header_env,omitempty"` + ResourceAttributes map[string]string `json:"resource_attributes,omitempty"` + ServiceName string `json:"service_name,omitempty"` + ServiceNamespace string `json:"service_namespace,omitempty"` + ServiceVersion string `json:"service_version,omitempty"` + InstrumentationScope string `json:"instrumentation_scope,omitempty"` + TimeoutMillis uint64 `json:"timeout_millis,omitempty"` + MaxQueueSize *uint64 `json:"max_queue_size,omitempty"` + MaxExportBatchSize *uint64 `json:"max_export_batch_size,omitempty"` + ScheduledDelayMillis *uint64 `json:"scheduled_delay_millis,omitempty"` + ActivationPolicy *ExportActivationPolicyConfig `json:"activation_policy,omitempty"` } // ObservabilityAtofConfig configures filesystem-backed raw ATOF JSONL export. @@ -98,9 +107,10 @@ type ObservabilityAtofSinkConfigurer interface { // ObservabilityAtofFileSinkConfig configures one filesystem ATOF JSONL destination. type ObservabilityAtofFileSinkConfig struct { - OutputDirectory string `json:"output_directory,omitempty"` - Filename string `json:"filename,omitempty"` - Mode string `json:"mode,omitempty"` + OutputDirectory string `json:"output_directory,omitempty"` + Filename string `json:"filename,omitempty"` + Mode string `json:"mode,omitempty"` + ActivationPolicy *ExportActivationPolicyConfig `json:"activation_policy,omitempty"` } func (ObservabilityAtofFileSinkConfig) atofSinkConfig() { @@ -118,13 +128,14 @@ func (config ObservabilityAtofFileSinkConfig) MarshalJSON() ([]byte, error) { // ObservabilityAtofStreamSinkConfig configures one remote ATOF destination. type ObservabilityAtofStreamSinkConfig struct { - URL string `json:"url"` - Transport string `json:"transport,omitempty"` - Headers map[string]string `json:"headers,omitempty"` - HeaderEnv map[string]string `json:"header_env,omitempty"` - TimeoutMillis uint64 `json:"timeout_millis,omitempty"` - FieldNamePolicy string `json:"field_name_policy,omitempty"` - Name string `json:"name,omitempty"` + URL string `json:"url"` + Transport string `json:"transport,omitempty"` + Headers map[string]string `json:"headers,omitempty"` + HeaderEnv map[string]string `json:"header_env,omitempty"` + TimeoutMillis uint64 `json:"timeout_millis,omitempty"` + FieldNamePolicy string `json:"field_name_policy,omitempty"` + Name string `json:"name,omitempty"` + ActivationPolicy *ExportActivationPolicyConfig `json:"activation_policy,omitempty"` } func (ObservabilityAtofStreamSinkConfig) atofSinkConfig() { @@ -146,15 +157,16 @@ type ObservabilityAtofEndpoint = ObservabilityAtofStreamSinkConfig // ObservabilityAtifConfig configures per-top-level-agent ATIF file export. type ObservabilityAtifConfig struct { - Enabled bool `json:"enabled,omitempty"` - AgentName string `json:"agent_name,omitempty"` - AgentVersion string `json:"agent_version,omitempty"` - ModelName string `json:"model_name,omitempty"` - ToolDefinitions []map[string]any `json:"tool_definitions,omitempty"` - Extra map[string]any `json:"extra,omitempty"` - OutputDirectory string `json:"output_directory,omitempty"` - FilenameTemplate string `json:"filename_template,omitempty"` - Storage []ObservabilityAtifStorageConfigurer `json:"storage,omitempty"` + Enabled bool `json:"enabled,omitempty"` + AgentName string `json:"agent_name,omitempty"` + AgentVersion string `json:"agent_version,omitempty"` + ModelName string `json:"model_name,omitempty"` + ToolDefinitions []map[string]any `json:"tool_definitions,omitempty"` + Extra map[string]any `json:"extra,omitempty"` + OutputDirectory string `json:"output_directory,omitempty"` + FilenameTemplate string `json:"filename_template,omitempty"` + Storage []ObservabilityAtifStorageConfigurer `json:"storage,omitempty"` + LocalActivationPolicy *ExportActivationPolicyConfig `json:"local_activation_policy,omitempty"` } // ObservabilityAtifStorageConfigurer is one remote ATIF trajectory storage destination. @@ -164,14 +176,15 @@ type ObservabilityAtifStorageConfigurer interface { // ObservabilityS3StorageConfig configures S3-compatible ATIF trajectory upload. type ObservabilityS3StorageConfig struct { - Bucket string `json:"bucket"` - KeyPrefix string `json:"key_prefix,omitempty"` - AccessKeyID string `json:"access_key_id,omitempty"` - SecretAccessKeyVar string `json:"secret_access_key_var,omitempty"` - SessionTokenVar string `json:"session_token_var,omitempty"` - Region string `json:"region,omitempty"` - EndpointURL string `json:"endpoint_url,omitempty"` - AllowHTTP *bool `json:"allow_http,omitempty"` + Bucket string `json:"bucket"` + KeyPrefix string `json:"key_prefix,omitempty"` + AccessKeyID string `json:"access_key_id,omitempty"` + SecretAccessKeyVar string `json:"secret_access_key_var,omitempty"` + SessionTokenVar string `json:"session_token_var,omitempty"` + Region string `json:"region,omitempty"` + EndpointURL string `json:"endpoint_url,omitempty"` + AllowHTTP *bool `json:"allow_http,omitempty"` + ActivationPolicy *ExportActivationPolicyConfig `json:"activation_policy,omitempty"` } func (ObservabilityS3StorageConfig) atifStorageConfig() { @@ -181,15 +194,16 @@ func (ObservabilityS3StorageConfig) atifStorageConfig() { // MarshalJSON serializes the S3 config with the core plugin's fixed type discriminator. func (config ObservabilityS3StorageConfig) MarshalJSON() ([]byte, error) { type s3StorageJSON struct { - Type string `json:"type"` - Bucket string `json:"bucket"` - KeyPrefix string `json:"key_prefix,omitempty"` - AccessKeyID string `json:"access_key_id,omitempty"` - SecretAccessKeyVar string `json:"secret_access_key_var,omitempty"` - SessionTokenVar string `json:"session_token_var,omitempty"` - Region string `json:"region,omitempty"` - EndpointURL string `json:"endpoint_url,omitempty"` - AllowHTTP *bool `json:"allow_http,omitempty"` + Type string `json:"type"` + Bucket string `json:"bucket"` + KeyPrefix string `json:"key_prefix,omitempty"` + AccessKeyID string `json:"access_key_id,omitempty"` + SecretAccessKeyVar string `json:"secret_access_key_var,omitempty"` + SessionTokenVar string `json:"session_token_var,omitempty"` + Region string `json:"region,omitempty"` + EndpointURL string `json:"endpoint_url,omitempty"` + AllowHTTP *bool `json:"allow_http,omitempty"` + ActivationPolicy *ExportActivationPolicyConfig `json:"activation_policy,omitempty"` } return json.Marshal(s3StorageJSON{ Type: "s3", @@ -201,15 +215,17 @@ func (config ObservabilityS3StorageConfig) MarshalJSON() ([]byte, error) { Region: config.Region, EndpointURL: config.EndpointURL, AllowHTTP: config.AllowHTTP, + ActivationPolicy: config.ActivationPolicy, }) } // ObservabilityHttpStorageConfig configures HTTP ATIF trajectory upload. type ObservabilityHttpStorageConfig struct { - Endpoint string `json:"endpoint"` - Headers map[string]string `json:"headers,omitempty"` - HeaderEnv map[string]string `json:"header_env,omitempty"` - TimeoutMillis uint64 `json:"timeout_millis,omitempty"` + Endpoint string `json:"endpoint"` + Headers map[string]string `json:"headers,omitempty"` + HeaderEnv map[string]string `json:"header_env,omitempty"` + TimeoutMillis uint64 `json:"timeout_millis,omitempty"` + ActivationPolicy *ExportActivationPolicyConfig `json:"activation_policy,omitempty"` } func (ObservabilityHttpStorageConfig) atifStorageConfig() { @@ -219,18 +235,20 @@ func (ObservabilityHttpStorageConfig) atifStorageConfig() { // MarshalJSON serializes the HTTP config with the core plugin's fixed type discriminator. func (config ObservabilityHttpStorageConfig) MarshalJSON() ([]byte, error) { type httpStorageJSON struct { - Type string `json:"type"` - Endpoint string `json:"endpoint"` - Headers map[string]string `json:"headers,omitempty"` - HeaderEnv map[string]string `json:"header_env,omitempty"` - TimeoutMillis uint64 `json:"timeout_millis,omitempty"` + Type string `json:"type"` + Endpoint string `json:"endpoint"` + Headers map[string]string `json:"headers,omitempty"` + HeaderEnv map[string]string `json:"header_env,omitempty"` + TimeoutMillis uint64 `json:"timeout_millis,omitempty"` + ActivationPolicy *ExportActivationPolicyConfig `json:"activation_policy,omitempty"` } return json.Marshal(httpStorageJSON{ - Type: "http", - Endpoint: config.Endpoint, - Headers: config.Headers, - HeaderEnv: config.HeaderEnv, - TimeoutMillis: config.TimeoutMillis, + Type: "http", + Endpoint: config.Endpoint, + Headers: config.Headers, + HeaderEnv: config.HeaderEnv, + TimeoutMillis: config.TimeoutMillis, + ActivationPolicy: config.ActivationPolicy, }) } @@ -245,6 +263,11 @@ func NewObservabilityConfig() ObservabilityConfig { return ObservabilityConfig{Version: 4} } +// NewExportActivationPolicyConfig returns an activation policy with the canonical timeout. +func NewExportActivationPolicyConfig(provider string) ExportActivationPolicyConfig { + return ExportActivationPolicyConfig{Provider: provider, TimeoutMillis: 30000} +} + // NewObservabilityAtofConfig returns disabled ATOF JSONL settings with native defaults. func NewObservabilityAtofConfig() ObservabilityAtofConfig { return ObservabilityAtofConfig{} diff --git a/go/nemo_relay/observability_plugin_test.go b/go/nemo_relay/observability_plugin_test.go index 90369ae69..0fcdc2eee 100644 --- a/go/nemo_relay/observability_plugin_test.go +++ b/go/nemo_relay/observability_plugin_test.go @@ -28,6 +28,9 @@ const ( ) func TestObservabilityConfigHelpers(t *testing.T) { + policyValue := NewExportActivationPolicyConfig("com.example.runtime-policy") + policyValue.Config = map[string]any{"allowed_countries": []string{"US", "CA"}} + policy := &policyValue config := NewObservabilityConfig() if config.Version != 4 { t.Fatalf("expected version 4, got %d", config.Version) @@ -40,18 +43,20 @@ func TestObservabilityConfigHelpers(t *testing.T) { t.Fatalf("unexpected ATOF defaults: %#v", atof) } atof.Sinks = []ObservabilityAtofSinkConfigurer{ObservabilityAtofEndpoint{ - Name: "archive", - URL: "http://localhost:8080/events", - Transport: "http_post", - Headers: map[string]string{"X-Test": "yes"}, - HeaderEnv: map[string]string{"authorization": "NEMO_RELAY_ATOF_AUTH"}, - TimeoutMillis: 1000, - FieldNamePolicy: "replace_dots", + Name: "archive", + URL: "http://localhost:8080/events", + Transport: "http_post", + Headers: map[string]string{"X-Test": "yes"}, + HeaderEnv: map[string]string{"authorization": "NEMO_RELAY_ATOF_AUTH"}, + TimeoutMillis: 1000, + FieldNamePolicy: "replace_dots", + ActivationPolicy: policy, }} atif := NewObservabilityAtifConfig() if atif.Enabled || atif.AgentName != "NeMo Relay" || atif.ModelName != "unknown" || atif.FilenameTemplate != "nemo-relay-atif-{session_id}.json" { t.Fatalf("unexpected ATIF defaults: %#v", atif) } + atif.LocalActivationPolicy = policy allowHTTP := false s3Storage := NewObservabilityS3StorageConfig("archive") s3Storage.KeyPrefix = "runs/" @@ -63,6 +68,8 @@ func TestObservabilityConfigHelpers(t *testing.T) { httpStorage.Headers = map[string]string{testStaticHeader: "value"} httpStorage.HeaderEnv = map[string]string{"authorization": "NEMO_RELAY_ATIF_HTTP_AUTH"} httpStorage.TimeoutMillis = 1500 + httpStorage.ActivationPolicy = policy + s3Storage.ActivationPolicy = policy assertS3StorageConfig(t, s3Storage) assertHTTPStorageConfig(t, httpStorage) atif.Storage = []ObservabilityAtifStorageConfigurer{ @@ -76,6 +83,7 @@ func TestObservabilityConfigHelpers(t *testing.T) { } otel.Endpoints[0].HeaderEnv["authorization"] = "OTEL_AUTHORIZATION" otel.Endpoints[0].PromoteMetadataPrefixes = []string{"nv."} + otel.Endpoints[0].ActivationPolicy = policy maxQueueSize := uint64(4096) maxExportBatchSize := uint64(256) scheduledDelayMillis := uint64(750) @@ -89,6 +97,7 @@ func TestObservabilityConfigHelpers(t *testing.T) { metricEndpoint := NewObservabilityOpenTelemetrySignalEndpointConfig("https://collector.example/custom/metrics") metricEndpoint.Headers["x-nv-project"] = "observability-dev" metricEndpoint.ResourceAttributes["nv.project"] = "observability-dev" + metricEndpoint.ActivationPolicy = policy metrics.Endpoints = ObservabilityOpenTelemetrySignalEndpoints(metricEndpoint) otel.Logs = &logs otel.Metrics = &metrics @@ -106,6 +115,8 @@ func TestObservabilityConfigHelpers(t *testing.T) { func TestObservabilityAtofSinkConfigConstructorsSerializeTheirDiscriminators(t *testing.T) { file := NewObservabilityAtofFileSinkConfig() + filePolicy := NewExportActivationPolicyConfig("com.example.runtime-policy") + file.ActivationPolicy = &filePolicy if file.Mode != "append" { t.Fatalf("file sink mode = %q, want append", file.Mode) } @@ -116,6 +127,9 @@ func TestObservabilityAtofSinkConfigConstructorsSerializeTheirDiscriminators(t * if !strings.Contains(string(fileJSON), `"type":"file"`) { t.Fatalf("file sink discriminator missing: %s", fileJSON) } + if !strings.Contains(string(fileJSON), `"activation_policy"`) { + t.Fatalf("file sink activation policy missing: %s", fileJSON) + } stream := NewObservabilityAtofStreamSinkConfig("http://localhost:8080/events") if stream.Transport != "http_post" || stream.TimeoutMillis != 3000 || stream.FieldNamePolicy != "preserve" { @@ -148,6 +162,7 @@ func assertWrappedObservabilityConfig(t *testing.T, wrapped PluginComponentSpec) firstSink["header_env"].(map[string]any)["authorization"] != "NEMO_RELAY_ATOF_AUTH" { t.Fatalf("expected serialized ATOF stream sink settings, got %#v", sinks) } + assertExportActivationPolicy(t, firstSink["activation_policy"]) serialized, err := json.Marshal(wrapped) if err != nil { t.Fatalf("marshal observability component failed: %v", err) @@ -163,6 +178,7 @@ func assertWrappedObservabilityConfig(t *testing.T, wrapped PluginComponentSpec) if otelEndpoints[0].(map[string]any)["header_env"].(map[string]any)["authorization"] != "OTEL_AUTHORIZATION" { t.Fatalf("expected OpenTelemetry header_env in serialized config: %#v", wrapped.Config) } + assertExportActivationPolicy(t, otelEndpoints[0].(map[string]any)["activation_policy"]) promotePrefixes := otelEndpoints[0].(map[string]any)["promote_metadata_prefixes"].([]any) if len(promotePrefixes) != 1 || promotePrefixes[0] != "nv." { t.Fatalf("expected OpenTelemetry metadata promotion prefixes in serialized config: %#v", wrapped.Config) @@ -189,6 +205,15 @@ func assertWrappedObservabilityConfig(t *testing.T, wrapped PluginComponentSpec) metricEndpoint["resource_attributes"].(map[string]any)["nv.project"] != "observability-dev" { t.Fatalf("expected OpenTelemetry metric settings in serialized config: %#v", metrics) } + assertExportActivationPolicy(t, metricEndpoint["activation_policy"]) +} + +func assertExportActivationPolicy(t *testing.T, value any) { + t.Helper() + policy, ok := value.(map[string]any) + if !ok || policy["provider"] != "com.example.runtime-policy" || policy["timeout_millis"] != float64(30000) { + t.Fatalf("unexpected export activation policy: %#v", value) + } } func TestObservabilitySignalEndpointOmittedVersusExplicitEmpty(t *testing.T) { @@ -296,6 +321,7 @@ func assertHTTPStorageConfig(t *testing.T, storage ObservabilityHttpStorageConfi func assertWrappedAtifStorageConfig(t *testing.T, atifConfig map[string]any) { t.Helper() + assertExportActivationPolicy(t, atifConfig["local_activation_policy"]) storage := atifConfig["storage"].([]any) if len(storage) != 2 { t.Fatalf("expected two ATIF storage destinations, got %#v", storage) @@ -412,9 +438,11 @@ func requireOtelRequest(t *testing.T, requests <-chan otelRequest, path, bodyFra t.Helper() timeout := time.NewTimer(5 * time.Second) defer timeout.Stop() + var observedPaths []string for { select { case request := <-requests: + observedPaths = append(observedPaths, request.Path) if request.Path != path { continue } @@ -423,7 +451,7 @@ func requireOtelRequest(t *testing.T, requests <-chan otelRequest, path, bodyFra } return case <-timeout.C: - t.Fatalf("timed out waiting for OTLP %s export", path) + t.Fatalf("timed out waiting for OTLP %s export; observed paths: %v", path, observedPaths) } } } diff --git a/go/nemo_relay/plugin.go b/go/nemo_relay/plugin.go index 1de3de731..4dfba8412 100644 --- a/go/nemo_relay/plugin.go +++ b/go/nemo_relay/plugin.go @@ -30,6 +30,8 @@ typedef char* (*NemoRelayLlmExecNextFn)(const char* native_json, void* next_ctx) typedef char* (*NemoRelayLlmExecInterceptCb)(void* user_data, const char* native_json, NemoRelayLlmExecNextFn next_fn, void* next_ctx); typedef char* (*NemoRelayToolExecNextFn)(const char* args_json, void* next_ctx); typedef char* (*NemoRelayToolExecInterceptCb)(void* user_data, const char* args_json, NemoRelayToolExecNextFn next_fn, void* next_ctx); +typedef char* (*NemoRelayExportActivationPolicyCb)(void* user_data, const char* request_json); +typedef int32_t (*NemoRelayExportTargetActivationCb)(void* user_data); extern int32_t nemo_relay_validate_plugin_config(const char* config_json, char** out_json); extern int32_t nemo_relay_initialize_plugins(const char* config_json, char** out_json); @@ -58,6 +60,8 @@ extern int32_t nemo_relay_plugin_context_register_tool_request_intercept(FfiPlug extern int32_t nemo_relay_plugin_context_register_llm_execution_intercept(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmExecInterceptCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_llm_stream_execution_intercept(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmExecInterceptCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_tool_execution_intercept(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayToolExecInterceptCb cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_plugin_context_register_export_activation_policy(FfiPluginContext* ctx, NemoRelayExportActivationPolicyCb cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_plugin_context_register_export_target(FfiPluginContext* ctx, const char* registration_json, NemoRelayExportTargetActivationCb cb, void* user_data, NemoRelayFreeFn free_fn); extern char* goPluginValidateTrampoline(void*, const char*); extern int32_t goPluginRegisterTrampoline(void*, const char*, FfiPluginContext*); @@ -72,10 +76,13 @@ extern char* goLlmConditionalTrampoline(void*, const void*); extern char* goLlmExecInterceptTrampoline(void*, const char*, NemoRelayLlmExecNextFn, void*); extern int32_t goLlmRequestInterceptTrampoline(void*, const char*, const void*, const char*, char**); extern char* goToolExecInterceptTrampoline(void*, const char*, NemoRelayToolExecNextFn, void*); +extern char* goExportActivationPolicyTrampoline(void*, const char*); +extern int32_t goExportTargetActivationTrampoline(void*); */ import "C" import ( + "encoding/json" "errors" "log" "runtime" @@ -289,6 +296,33 @@ type PluginContext struct { ptr *C.FfiPluginContext } +// ExportActivationDecision is the activation-time decision for one export target. +type ExportActivationDecision string + +const ( + ExportActivationAllow ExportActivationDecision = "allow" + ExportActivationDeny ExportActivationDecision = "deny" +) + +// ExportActivationRequest is passed to a policy without exporter destinations or credentials. +type ExportActivationRequest struct { + TargetKind string `json:"target_kind"` + Config json.RawMessage `json:"config"` +} + +// ExportTargetRegistration describes one deferred plugin-managed exporter. +type ExportTargetRegistration struct { + ID string `json:"id"` + TargetKind string `json:"target_kind"` + ActivationPolicy *ExportActivationPolicyConfig `json:"activation_policy,omitempty"` +} + +// ExportActivationPolicyFunc decides whether Relay may activate one exporter. +type ExportActivationPolicyFunc func(ExportActivationRequest) (ExportActivationDecision, error) + +// ExportTargetActivationFunc constructs and starts an allowed exporter. +type ExportTargetActivationFunc func() error + // Plugin is the plugin callback contract. // // Validate receives one component-local config object and returns diagnostics. @@ -583,6 +617,41 @@ func (ctx *PluginContext) RegisterSubscriber(name string, fn EventSubscriberFunc )) } +// RegisterExportActivationPolicy registers the policy owned by this plugin kind. +func (ctx *PluginContext) RegisterExportActivationPolicy(fn ExportActivationPolicyFunc) error { + if ctx == nil || ctx.ptr == nil { + return errors.New(errPluginContextClosed) + } + userData := registerClosure(fn) + return checkStatus(C.nemo_relay_plugin_context_register_export_activation_policy( + ctx.ptr, + (C.NemoRelayExportActivationPolicyCb)(C.goExportActivationPolicyTrampoline), + userData, + (C.NemoRelayFreeFn)(C.goFreeTrampoline), + )) +} + +// RegisterExportTarget defers exporter construction until its policy allows activation. +func (ctx *PluginContext) RegisterExportTarget(registration ExportTargetRegistration, fn ExportTargetActivationFunc) error { + if ctx == nil || ctx.ptr == nil { + return errors.New(errPluginContextClosed) + } + payload, err := jsonMarshal(registration) + if err != nil { + return err + } + cPayload := C.CString(string(payload)) + defer C.free(unsafe.Pointer(cPayload)) + userData := registerClosure(fn) + return checkStatus(C.nemo_relay_plugin_context_register_export_target( + ctx.ptr, + cPayload, + (C.NemoRelayExportTargetActivationCb)(C.goExportTargetActivationTrampoline), + userData, + (C.NemoRelayFreeFn)(C.goFreeTrampoline), + )) +} + func (ctx *PluginContext) registerEventSanitizer(name string, priority int32, fn EventSanitizeFunc, surface int) error { if ctx == nil || ctx.ptr == nil { return errors.New(errPluginContextClosed) diff --git a/go/nemo_relay/plugin_gap_test.go b/go/nemo_relay/plugin_gap_test.go index 856b6080f..3d2355a60 100644 --- a/go/nemo_relay/plugin_gap_test.go +++ b/go/nemo_relay/plugin_gap_test.go @@ -5,6 +5,7 @@ package nemo_relay import ( "encoding/json" + "fmt" "testing" ) @@ -35,6 +36,66 @@ func TestPluginConfigSerializationErrorsSurfaceBeforeFFI(t *testing.T) { } } +func TestPluginCanGateItsOwnExportTarget(t *testing.T) { + for _, decision := range []ExportActivationDecision{ExportActivationAllow, ExportActivationDeny} { + t.Run(string(decision), func(t *testing.T) { + kind := fmt.Sprintf("tests.go_export_activation_%s", decision) + activations := 0 + err := RegisterPlugin(kind, PluginFuncs{ + RegisterFunc: func(_ map[string]any, ctx *PluginContext) error { + if err := ctx.RegisterExportActivationPolicy(func(request ExportActivationRequest) (ExportActivationDecision, error) { + if request.TargetKind != "tests.telemetry.otlp" { + t.Fatalf("unexpected target kind %q", request.TargetKind) + } + return decision, nil + }); err != nil { + return err + } + return ctx.RegisterExportTarget( + ExportTargetRegistration{ + ID: "self-otel", + TargetKind: "tests.telemetry.otlp", + ActivationPolicy: &ExportActivationPolicyConfig{ + Provider: kind, + TimeoutMillis: 30_000, + Config: json.RawMessage(`{"country":"US"}`), + }, + }, + func() error { + activations++ + return nil + }, + ) + }, + }) + if err != nil { + t.Fatalf("RegisterPlugin() error = %v", err) + } + defer DeregisterPlugin(kind) + defer ClearPluginConfiguration() + + _, err = InitializePlugins(PluginConfig{ + Version: 1, + Components: []PluginComponentSpec{{ + Kind: kind, + Enabled: true, + Config: map[string]any{}, + }}, + }) + if err != nil { + t.Fatalf("InitializePlugins() error = %v", err) + } + expected := 0 + if decision == ExportActivationAllow { + expected = 1 + } + if activations != expected { + t.Fatalf("activation count = %d, want %d", activations, expected) + } + }) + } +} + func TestClosedPluginContextRejectsEveryRegistrationSurface(t *testing.T) { ctx := &PluginContext{} request := func(LLMRequestDTO, LLMSanitizeRequestContext) (LLMRequestDTO, bool) { @@ -49,6 +110,12 @@ func TestClosedPluginContextRejectsEveryRegistrationSurface(t *testing.T) { call func() error }{ {name: "subscriber", call: func() error { return ctx.RegisterSubscriber("closed_subscriber", nil) }}, + {name: "export activation policy", call: func() error { + return ctx.RegisterExportActivationPolicy(nil) + }}, + {name: "export target", call: func() error { + return ctx.RegisterExportTarget(ExportTargetRegistration{}, nil) + }}, {name: "mark sanitizer", call: func() error { return ctx.RegisterMarkSanitizeGuardrail("closed_mark", 0, nil) }}, {name: "scope-start sanitizer", call: func() error { return ctx.RegisterScopeSanitizeStartGuardrail("closed_scope_start", 0, nil) }}, {name: "scope-end sanitizer", call: func() error { return ctx.RegisterScopeSanitizeEndGuardrail("closed_scope_end", 0, nil) }}, diff --git a/justfile b/justfile index 9de25f006..3da955bed 100644 --- a/justfile +++ b/justfile @@ -1058,10 +1058,12 @@ check-python-worker-proto: assert pb.HandshakeRequest.DESCRIPTOR.fields_by_name["worker_protocol"].number == 4 assert pb.InvokeRequest.DESCRIPTOR.fields_by_name["auth_token"].number == 7 + assert pb.InvokeRequest.DESCRIPTOR.fields_by_name["export_activation"].number == 13 assert {method.name for method in pb.DESCRIPTOR.services_by_name["PluginWorker"].methods} == { "Handshake", "Health", "Validate", "Register", "Invoke", "InvokeStream", "CancelInvocation", "Shutdown" } assert pb.SUBSCRIBER == 1 + assert pb.EXPORT_ACTIVATION_POLICY == 3 assert pb.LLM_STREAM_EXECUTION_INTERCEPT == 25 tool_next = pb.DESCRIPTOR.services_by_name["RelayHostRuntime"].methods_by_name["ToolNext"] assert tool_next.output_type.full_name == "nemo.relay.worker.v1.ToolExecutionResultResponse" diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index ee4a46ae4..bfe0504c3 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -1451,6 +1451,8 @@ class PluginContext: Python plugin protocols expose the public shape. The native class exists for runtime registration callbacks. """ + def register_export_activation_policy(self, callback: Callable[[object], object]) -> None: ... + def register_export_target(self, registration: object, callback: Callable[[], object]) -> None: ... def register_mark_sanitize_guardrail(self, name: str, priority: int, callback: _EventSanitizeGuardrail) -> None: ... def register_scope_sanitize_start_guardrail( self, name: str, priority: int, callback: _EventSanitizeGuardrail diff --git a/python/nemo_relay/observability.py b/python/nemo_relay/observability.py index 113ef1575..6a7365c88 100644 --- a/python/nemo_relay/observability.py +++ b/python/nemo_relay/observability.py @@ -52,6 +52,23 @@ def to_dict(self) -> JsonObject: } +@dataclass(slots=True) +class ExportActivationPolicyConfig: + """Activation-time policy attached to one export target.""" + + provider: str + timeout_millis: int = 30000 + config: Json = None + + def to_dict(self) -> JsonObject: + """Serialize this policy attachment to the canonical plugin shape.""" + return { + "provider": self.provider, + "timeout_millis": self.timeout_millis, + "config": self.config, + } + + @dataclass(slots=True) class AtofStreamSinkConfig: """Stream sink for raw ATOF events.""" @@ -63,6 +80,7 @@ class AtofStreamSinkConfig: timeout_millis: int = 3000 field_name_policy: Literal["preserve", "replace_dots"] = "preserve" name: str | None = None + activation_policy: ExportActivationPolicyConfig | None = None def to_dict(self) -> JsonObject: """Serialize this ATOF stream sink to the canonical JSON object shape.""" @@ -76,6 +94,7 @@ def to_dict(self) -> JsonObject: "header_env": self.header_env, "timeout_millis": self.timeout_millis, "field_name_policy": self.field_name_policy, + "activation_policy": self.activation_policy, } ) @@ -104,6 +123,7 @@ class AtofFileSinkConfig: output_directory: str | None = None filename: str | None = None mode: Literal["append", "overwrite"] = "append" + activation_policy: ExportActivationPolicyConfig | None = None def to_dict(self) -> JsonObject: return _normalize_object( @@ -112,6 +132,7 @@ def to_dict(self) -> JsonObject: "output_directory": self.output_directory, "filename": self.filename, "mode": self.mode, + "activation_policy": self.activation_policy, } ) @@ -139,6 +160,7 @@ class S3StorageConfig: region: str | None = None endpoint_url: str | None = None allow_http: bool | None = None + activation_policy: ExportActivationPolicyConfig | None = None def to_dict(self) -> JsonObject: """Serialize this S3 storage config to the canonical JSON object shape.""" @@ -153,6 +175,7 @@ def to_dict(self) -> JsonObject: "region": self.region, "endpoint_url": self.endpoint_url, "allow_http": self.allow_http, + "activation_policy": self.activation_policy, } ) @@ -165,6 +188,7 @@ class HttpStorageConfig: headers: dict[str, str] = field(default_factory=dict) header_env: dict[str, str] = field(default_factory=dict) timeout_millis: int = 3000 + activation_policy: ExportActivationPolicyConfig | None = None def to_dict(self) -> JsonObject: """Serialize this HTTP storage config to the canonical JSON object shape.""" @@ -175,6 +199,7 @@ def to_dict(self) -> JsonObject: "headers": self.headers, "header_env": self.header_env, "timeout_millis": self.timeout_millis, + "activation_policy": self.activation_policy, } ) @@ -192,6 +217,7 @@ class AtifConfig: output_directory: str | None = None filename_template: str = "nemo-relay-atif-{session_id}.json" storage: list[S3StorageConfig | HttpStorageConfig] | None = None + local_activation_policy: ExportActivationPolicyConfig | None = None def to_dict(self) -> JsonObject: """Serialize this ATIF config to the canonical JSON object shape.""" @@ -205,6 +231,7 @@ def to_dict(self) -> JsonObject: "output_directory": self.output_directory, "filename_template": self.filename_template, "storage": self.storage, + "local_activation_policy": self.local_activation_policy, } if value["agent_version"] is None: value.pop("agent_version") @@ -233,6 +260,7 @@ class OpenTelemetryEndpointConfig: max_export_batch_size: int | None = None scheduled_delay_millis: int | None = None promote_metadata_prefixes: list[str] = field(default_factory=list) + activation_policy: ExportActivationPolicyConfig | None = None def to_dict(self) -> JsonObject: """Serialize this endpoint to the canonical plugin shape.""" @@ -256,6 +284,7 @@ def to_dict(self) -> JsonObject: "headers": self.headers, "header_env": self.header_env, "resource_attributes": self.resource_attributes, + "activation_policy": self.activation_policy, } ) @@ -274,6 +303,7 @@ class OpenTelemetrySignalEndpointConfig: service_version: str | None = None instrumentation_scope: str = "opentelemetry" timeout_millis: int = 3000 + activation_policy: ExportActivationPolicyConfig | None = None def to_dict(self) -> JsonObject: """Serialize this signal endpoint to the canonical plugin shape.""" @@ -289,6 +319,7 @@ def to_dict(self) -> JsonObject: "service_version": self.service_version, "instrumentation_scope": self.instrumentation_scope, "timeout_millis": self.timeout_millis, + "activation_policy": self.activation_policy, } ) @@ -414,6 +445,7 @@ def to_dict(self) -> JsonObject: __all__ = [ "ConfigPolicy", + "ExportActivationPolicyConfig", "AtofEndpointConfig", "AtofFileSinkConfig", "AtofStreamSinkConfig", diff --git a/python/nemo_relay/observability.pyi b/python/nemo_relay/observability.pyi index b8d358cc3..447272188 100644 --- a/python/nemo_relay/observability.pyi +++ b/python/nemo_relay/observability.pyi @@ -8,7 +8,7 @@ from __future__ import annotations from dataclasses import dataclass, field from typing import Literal -from nemo_relay import JsonObject, UnsupportedBehavior +from nemo_relay import Json, JsonObject, UnsupportedBehavior @dataclass(slots=True) class ConfigPolicy: @@ -17,6 +17,13 @@ class ConfigPolicy: unsupported_value: UnsupportedBehavior = ... def to_dict(self) -> JsonObject: ... +@dataclass(slots=True) +class ExportActivationPolicyConfig: + provider: str + timeout_millis: int = ... + config: Json = ... + def to_dict(self) -> JsonObject: ... + @dataclass(slots=True) class AtofStreamSinkConfig: url: str = ... @@ -26,6 +33,7 @@ class AtofStreamSinkConfig: timeout_millis: int = ... field_name_policy: Literal["preserve", "replace_dots"] = ... name: str | None = ... + activation_policy: ExportActivationPolicyConfig | None = ... def to_dict(self) -> JsonObject: ... @dataclass(slots=True) @@ -39,6 +47,7 @@ class AtofFileSinkConfig: output_directory: str | None = ... filename: str | None = ... mode: Literal["append", "overwrite"] = ... + activation_policy: ExportActivationPolicyConfig | None = ... def to_dict(self) -> JsonObject: ... AtofEndpointConfig = AtofStreamSinkConfig @@ -53,6 +62,7 @@ class S3StorageConfig: region: str | None = ... endpoint_url: str | None = ... allow_http: bool | None = ... + activation_policy: ExportActivationPolicyConfig | None = ... def to_dict(self) -> JsonObject: ... @dataclass(slots=True) @@ -61,6 +71,7 @@ class HttpStorageConfig: headers: dict[str, str] = field(default_factory=dict) header_env: dict[str, str] = field(default_factory=dict) timeout_millis: int = ... + activation_policy: ExportActivationPolicyConfig | None = ... def to_dict(self) -> JsonObject: ... @dataclass(slots=True) @@ -74,6 +85,7 @@ class AtifConfig: output_directory: str | None = ... filename_template: str = ... storage: list[S3StorageConfig | HttpStorageConfig] | None = ... + local_activation_policy: ExportActivationPolicyConfig | None = ... def to_dict(self) -> JsonObject: ... @dataclass(slots=True) @@ -96,6 +108,7 @@ class OpenTelemetryEndpointConfig: max_export_batch_size: int | None = ... scheduled_delay_millis: int | None = ... promote_metadata_prefixes: list[str] = field(default_factory=list) + activation_policy: ExportActivationPolicyConfig | None = ... def to_dict(self) -> JsonObject: ... @dataclass(slots=True) @@ -110,6 +123,7 @@ class OpenTelemetrySignalEndpointConfig: service_version: str | None = ... instrumentation_scope: str = ... timeout_millis: int = ... + activation_policy: ExportActivationPolicyConfig | None = ... def to_dict(self) -> JsonObject: ... @dataclass(slots=True) diff --git a/python/nemo_relay/plugin.py b/python/nemo_relay/plugin.py index 6afe38b64..f2eb776d1 100644 --- a/python/nemo_relay/plugin.py +++ b/python/nemo_relay/plugin.py @@ -18,7 +18,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass, field, fields, is_dataclass from pathlib import Path -from typing import TYPE_CHECKING, AsyncIterator, Callable, Literal, Protocol, Self, TypedDict, cast +from typing import TYPE_CHECKING, AsyncIterator, Awaitable, Callable, Literal, Protocol, Self, TypedDict, cast from nemo_relay import ( EventSanitizeGuardrail, @@ -110,9 +110,50 @@ class ConfigReport(TypedDict): """Execution lane for a dynamically loaded plugin.""" +class ExportActivationPolicyConfig(TypedDict, total=False): + """Policy provider and bounded evaluation settings for one exporter.""" + + provider: str + timeout_millis: int + config: Json + + +class ExportActivationRequest(TypedDict): + """Secret-free request passed to an exporter activation policy.""" + + target_kind: str + config: Json + + +class ExportTargetRegistration(TypedDict, total=False): + """Deferred plugin-managed exporter registration.""" + + id: str + target_kind: str + activation_policy: ExportActivationPolicyConfig + + class PluginContext(Protocol): """Component-scoped registration context passed to custom plugin handlers.""" + def register_export_activation_policy( + self, + callback: Callable[ + [ExportActivationRequest], + Literal["allow", "deny"] | Awaitable[Literal["allow", "deny"]], + ], + ) -> None: + """Register the activation policy owned by this plugin kind.""" + ... + + def register_export_target( + self, + registration: ExportTargetRegistration, + callback: Callable[[], None | Awaitable[None]], + ) -> None: + """Defer exporter construction until activation allows the target.""" + ... + def register_subscriber(self, name: str, callback: Callable[[Event], None]) -> None: """Register an infallible event subscriber for this component.""" ... @@ -714,6 +755,9 @@ def deregister(plugin_kind: str) -> bool: "ConfigReport", "DynamicPluginActivationSpec", "DynamicPluginKind", + "ExportActivationPolicyConfig", + "ExportActivationRequest", + "ExportTargetRegistration", "PluginConfig", "PluginContext", "PluginHostActivation", diff --git a/python/nemo_relay/plugin.pyi b/python/nemo_relay/plugin.pyi index ee8f165dd..fd953c0a2 100644 --- a/python/nemo_relay/plugin.pyi +++ b/python/nemo_relay/plugin.pyi @@ -4,11 +4,12 @@ import os from collections.abc import Callable, Sequence from types import TracebackType -from typing import AsyncContextManager, Literal, Protocol, Self, TypedDict +from typing import AsyncContextManager, Awaitable, Literal, Protocol, Self, TypedDict from nemo_relay import ( Event, EventSanitizeGuardrail, + Json, JsonObject, LlmConditionalExecutionGuardrail, LlmExecutionIntercept, @@ -48,7 +49,28 @@ class ConfigReport(TypedDict): diagnostics: list[ConfigDiagnostic] runtime_diagnostics: list[RuntimeDiagnostic] +class ExportActivationPolicyConfig(TypedDict, total=False): + provider: str + timeout_millis: int + config: Json + +class ExportActivationRequest(TypedDict): + target_kind: str + config: Json + +class ExportTargetRegistration(TypedDict, total=False): + id: str + target_kind: str + activation_policy: ExportActivationPolicyConfig + class PluginContext(Protocol): + def register_export_activation_policy( + self, + callback: Callable[[ExportActivationRequest], Literal["allow", "deny"] | Awaitable[Literal["allow", "deny"]]], + ) -> None: ... + def register_export_target( + self, registration: ExportTargetRegistration, callback: Callable[[], None | Awaitable[None]] + ) -> None: ... def register_subscriber(self, name: str, callback: Callable[[Event], None]) -> None: ... def register_mark_sanitize_guardrail(self, name: str, priority: int, callback: EventSanitizeGuardrail) -> None: ... def register_scope_sanitize_start_guardrail( diff --git a/python/plugin/README.md b/python/plugin/README.md index 04e001eaf..02c851409 100644 --- a/python/plugin/README.md +++ b/python/plugin/README.md @@ -35,6 +35,8 @@ regenerate their protobuf bindings, rebuild with this SDK, and declare a code outside the Relay host process. - **Use the shared runtime contract**: Register subscribers, guardrails, and intercepts through `WorkerPlugin` and `PluginContext`. +- **Gate exporters at activation**: Register one typed export policy + that receives only the target kind and opaque policy configuration. - **Call back into Relay safely**: Emit marks, create scopes, and continue managed execution through the host runtime handle. - **Keep worker lifecycle managed**: Let Relay provision the worker environment, diff --git a/python/plugin/src/nemo_relay_plugin/__init__.py b/python/plugin/src/nemo_relay_plugin/__init__.py index 2f8b8f588..ba34750d7 100644 --- a/python/plugin/src/nemo_relay_plugin/__init__.py +++ b/python/plugin/src/nemo_relay_plugin/__init__.py @@ -45,6 +45,11 @@ ToolExecutionInterceptOutcome: Canonical tool execution-intercept result. DiagnosticLevel: Severity of a configuration diagnostic. ConfigDiagnostic: Structured configuration warning or error. + ExportActivationTargetKind: Local or remote exporter kind presented to a policy. + ExportActivationRequest: Secret-free activation request for one export target. + ExportActivationDecision: Allow-or-deny result returned by an activation policy. + ExportActivationPolicyConfig: Provider and timeout settings attached to an export target. + ExportTargetRegistration: Deferred plugin-owned exporter registration. RuntimeDiagnostic: One aggregated host runtime diagnostic entry. RuntimeDiagnostics: Immutable host runtime diagnostics snapshot. ScopeType: Semantic category for a Relay execution scope. @@ -53,6 +58,8 @@ Public callback aliases: SubscriberCallback: Event subscriber callback. EventMetadataInjectorCallback: Event metadata injector callback. + ExportActivationPolicyCallback: Activation-time export policy callback. + ExportTargetActivationCallback: Callback that constructs an allowed exporter. EventSanitizeCallback: Mark or scope event sanitizer callback. ToolSanitizeCallback: Tool request or response sanitizer callback. ToolConditionalCallback: Tool execution guardrail callback. @@ -86,6 +93,13 @@ EventMetadataInjectorCallback, EventSanitizeCallback, EventSanitizeFields, + ExportActivationDecision, + ExportActivationPolicyCallback, + ExportActivationPolicyConfig, + ExportActivationRequest, + ExportActivationTargetKind, + ExportTargetActivationCallback, + ExportTargetRegistration, Json, LlmCodecIdentity, LlmConditionalCallback, @@ -141,6 +155,13 @@ "EventMetadataInjectorCallback", "EventSanitizeCallback", "EventSanitizeFields", + "ExportActivationDecision", + "ExportActivationPolicyCallback", + "ExportActivationPolicyConfig", + "ExportActivationRequest", + "ExportActivationTargetKind", + "ExportTargetActivationCallback", + "ExportTargetRegistration", "Json", "LlmConditionalCallback", "LlmCodecIdentity", diff --git a/python/plugin/src/nemo_relay_plugin/_api.py b/python/plugin/src/nemo_relay_plugin/_api.py index e10de0cbe..776e3962d 100644 --- a/python/plugin/src/nemo_relay_plugin/_api.py +++ b/python/plugin/src/nemo_relay_plugin/_api.py @@ -27,6 +27,9 @@ LlmRequestInterceptOutcome: Canonical LLM request-intercept result. DiagnosticLevel: Severity of a configuration diagnostic. ConfigDiagnostic: Structured configuration warning or error. + ExportActivationTargetKind: Local or remote exporter kind presented to a policy. + ExportActivationRequest: Secret-free activation request for one export target. + ExportActivationDecision: Allow-or-deny result returned by an activation policy. ScopeType: Semantic category for a Relay execution scope. WorkerSdkError: SDK, host-call, or worker protocol error. @@ -41,6 +44,7 @@ Public callback aliases used in registration annotations: SubscriberCallback: Event subscriber callback. EventMetadataInjectorCallback: Event metadata injector callback. + ExportActivationPolicyCallback: Activation-time export policy callback. ToolSanitizeCallback: Tool request or response sanitizer callback. ToolConditionalCallback: Tool execution guardrail callback. ToolRequestCallback: Tool request intercept callback. @@ -226,6 +230,7 @@ def _llm_codec_capability(invocation: pb.LlmInvocation) -> str | None: WORKER_PROTOCOL = "grpc-v1" JSON_SCHEMA = "nemo.relay.Json@1" +EXPORT_ACTIVATION_REQUEST_SCHEMA = "nemo.relay.ExportActivationRequest@1" DATA_SCHEMA_SCHEMA = "nemo.relay.DataSchema@1" METRIC_DATA_SCHEMA_NAME = "nemo.relay.metric_measurements" METRIC_DATA_SCHEMA_VERSION = "1" @@ -937,6 +942,82 @@ def register(self, ctx: PluginContext, config: Json) -> None | Awaitable[None]: SubscriberCallback: TypeAlias = Callable[[Event], None | Awaitable[None]] EventMetadataInjectorCallback: TypeAlias = Callable[[Event], dict[str, Json] | Awaitable[dict[str, Json]]] + + +class ExportActivationTargetKind(str): + """Validated namespaced exporter kind presented to an activation policy.""" + + OTLP_TRACE: ClassVar[ExportActivationTargetKind] + OTLP_LOG: ClassVar[ExportActivationTargetKind] + OTLP_METRIC: ClassVar[ExportActivationTargetKind] + ATOF_FILE: ClassVar[ExportActivationTargetKind] + ATOF_STREAM: ClassVar[ExportActivationTargetKind] + ATIF_FILE: ClassVar[ExportActivationTargetKind] + ATIF_HTTP: ClassVar[ExportActivationTargetKind] + ATIF_S3: ClassVar[ExportActivationTargetKind] + + def __new__(cls, value: str) -> ExportActivationTargetKind: + if ( + not isinstance(value, str) + or len(value) > 255 + or "." not in value + or any( + not segment or not all(character.isalnum() or character in "_-" for character in segment) + for segment in value.split(".") + ) + ): + raise ValueError("export target kind must be a dot-separated namespaced identifier") + return str.__new__(cls, value) + + +ExportActivationTargetKind.OTLP_TRACE = ExportActivationTargetKind("nemo_relay.otlp.trace") +ExportActivationTargetKind.OTLP_LOG = ExportActivationTargetKind("nemo_relay.otlp.log") +ExportActivationTargetKind.OTLP_METRIC = ExportActivationTargetKind("nemo_relay.otlp.metric") +ExportActivationTargetKind.ATOF_FILE = ExportActivationTargetKind("nemo_relay.atof.file") +ExportActivationTargetKind.ATOF_STREAM = ExportActivationTargetKind("nemo_relay.atof.stream") +ExportActivationTargetKind.ATIF_FILE = ExportActivationTargetKind("nemo_relay.atif.file") +ExportActivationTargetKind.ATIF_HTTP = ExportActivationTargetKind("nemo_relay.atif.http") +ExportActivationTargetKind.ATIF_S3 = ExportActivationTargetKind("nemo_relay.atif.s3") + + +class ExportActivationDecision(str, Enum): + """Decision returned by an export-activation policy.""" + + ALLOW = "allow" + DENY = "deny" + + +@dataclass(frozen=True, slots=True) +class ExportActivationRequest: + """Non-secret input supplied to an export-activation policy.""" + + target_kind: ExportActivationTargetKind + config: Json = None + + +@dataclass(frozen=True, slots=True) +class ExportActivationPolicyConfig: + """Policy provider and bounded evaluation configuration for one target.""" + + provider: str + timeout_millis: int = 30_000 + config: Json = None + + +@dataclass(frozen=True, slots=True) +class ExportTargetRegistration: + """Deferred plugin-managed exporter registration.""" + + id: str + target_kind: ExportActivationTargetKind + activation_policy: ExportActivationPolicyConfig | None = None + + +ExportActivationPolicyCallback: TypeAlias = Callable[ + [ExportActivationRequest], + ExportActivationDecision | Awaitable[ExportActivationDecision], +] +ExportTargetActivationCallback: TypeAlias = Callable[[], None | Awaitable[None]] EventSanitizeCallback: TypeAlias = Callable[ [Event, EventSanitizeFields], EventSanitizeFields | Awaitable[EventSanitizeFields], @@ -971,6 +1052,8 @@ class _Handlers: registrations: list[Any] subscribers: dict[str, SubscriberCallback] event_metadata_injectors: dict[str, EventMetadataInjectorCallback] + export_activation_policies: dict[str, ExportActivationPolicyCallback] + export_targets: dict[str, ExportTargetActivationCallback] mark_sanitizers: dict[str, EventSanitizeCallback] scope_start_sanitizers: dict[str, EventSanitizeCallback] scope_end_sanitizers: dict[str, EventSanitizeCallback] @@ -992,6 +1075,8 @@ def empty(cls) -> _Handlers: registrations=[], subscribers={}, event_metadata_injectors={}, + export_activation_policies={}, + export_targets={}, mark_sanitizers={}, scope_start_sanitizers={}, scope_end_sanitizers={}, @@ -1097,6 +1182,29 @@ def register_event_metadata_injector( self._push_registration(name, pb.EVENT_METADATA_INJECTOR, priority, False) self._handlers.event_metadata_injectors[name] = callback + def register_export_activation_policy(self, callback: ExportActivationPolicyCallback) -> None: + """Register this plugin's activation-time policy for exporters.""" + name = "export_activation_policy" + self._push_registration(name, pb.EXPORT_ACTIVATION_POLICY, 0, False) + self._handlers.export_activation_policies[name] = callback + + def register_export_target( + self, + registration: ExportTargetRegistration, + callback: ExportTargetActivationCallback, + ) -> None: + """Register a callback that constructs and starts an allowed exporter.""" + metadata = asdict(registration) + metadata["target_kind"] = str(registration.target_kind) + self._push_registration( + registration.id, + pb.EXPORT_TARGET, + 0, + False, + export_target=_json_envelope("nemo.relay.ExportTargetRegistration@1", metadata), + ) + self._handlers.export_targets[registration.id] = callback + def _register_event_sanitizer( self, name: str, @@ -1404,7 +1512,15 @@ def register_llm_stream_execution_intercept( self._push_registration(name, pb.LLM_STREAM_EXECUTION_INTERCEPT, priority, False) self._handlers.llm_stream_executions[name] = callback - def _push_registration(self, name: str, surface: int, priority: int, break_chain: bool) -> None: + def _push_registration( + self, + name: str, + surface: int, + priority: int, + break_chain: bool, + *, + export_target: Any | None = None, + ) -> None: if any( registration.local_name == name and registration.surface == surface for registration in self._handlers.registrations @@ -1416,6 +1532,7 @@ def _push_registration(self, name: str, surface: int, priority: int, break_chain surface=surface, priority=priority, break_chain=break_chain, + export_target=export_target, ) ) @@ -2230,6 +2347,33 @@ async def _invoke_result(self, request: Any) -> Any: )(event) ) return _json_response(result) + if request.surface == pb.EXPORT_ACTIVATION_POLICY: + payload = _decode_required_envelope( + request.export_activation, + "export activation request", + EXPORT_ACTIVATION_REQUEST_SCHEMA, + ) + if not isinstance(payload, dict): + raise WorkerSdkError("export activation request must be an object") + try: + policy_request = ExportActivationRequest( + target_kind=ExportActivationTargetKind(payload["target_kind"]), + config=payload.get("config"), + ) + except (KeyError, TypeError, ValueError) as exc: + raise WorkerSdkError(f"invalid export activation request: {exc}") from exc + decision = await _maybe_await( + self._handler( + self._handlers.export_activation_policies, + request.registration_name, + )(policy_request) + ) + if not isinstance(decision, ExportActivationDecision): + raise WorkerSdkError("export activation policy must return ExportActivationDecision") + return _json_response(decision.value) + if request.surface == pb.EXPORT_TARGET: + await _maybe_await(self._handler(self._handlers.export_targets, request.registration_name)()) + return pb.InvokeResponse(empty=pb.EmptyResult()) if request.surface in PluginContext._EVENT_SANITIZER_HANDLER_ATTRIBUTES: event = _decode_required_envelope(request.event, "event", EVENT_SCHEMA) fields: EventSanitizeFields = { @@ -2414,6 +2558,8 @@ def _all_surfaces() -> list[int]: return [ pb.SUBSCRIBER, pb.EVENT_METADATA_INJECTOR, + pb.EXPORT_ACTIVATION_POLICY, + pb.EXPORT_TARGET, pb.MARK_SANITIZE_GUARDRAIL, pb.SCOPE_SANITIZE_START_GUARDRAIL, pb.SCOPE_SANITIZE_END_GUARDRAIL, diff --git a/python/tests/plugin/test_public_api_docstrings.py b/python/tests/plugin/test_public_api_docstrings.py index 166539858..0df03409a 100644 --- a/python/tests/plugin/test_public_api_docstrings.py +++ b/python/tests/plugin/test_public_api_docstrings.py @@ -23,6 +23,8 @@ "Event", "EventMetadataInjectorCallback", "EventSanitizeCallback", + "ExportActivationPolicyCallback", + "ExportTargetActivationCallback", "Json", "LlmRequest", "SubscriberCallback", diff --git a/python/tests/plugin/test_python_worker_example.py b/python/tests/plugin/test_python_worker_example.py index 70bffb42f..a79f5b298 100644 --- a/python/tests/plugin/test_python_worker_example.py +++ b/python/tests/plugin/test_python_worker_example.py @@ -23,7 +23,13 @@ pytest.importorskip("grpc") -from nemo_relay_plugin import PluginContext, PluginRuntime # noqa: E402 +from nemo_relay_plugin import ( # noqa: E402 + ExportActivationDecision, + ExportActivationRequest, + ExportActivationTargetKind, + PluginContext, + PluginRuntime, +) def test_manifest_integrity_matches_artifact_bytes(): @@ -126,6 +132,19 @@ async def test_example_register_propagates_configured_tag(example: Any): plugin = example.ExamplePythonWorker() plugin.register(context, {"tag": "demo"}) + context.register_export_activation_policy.assert_called_once() + (decide_export,) = context.register_export_activation_policy.call_args.args + request = ExportActivationRequest( + target_kind=ExportActivationTargetKind.OTLP_TRACE, + config={"allow": True}, + ) + assert await decide_export(request) is ExportActivationDecision.ALLOW + request = ExportActivationRequest( + target_kind=ExportActivationTargetKind.OTLP_TRACE, + config={"allow": False}, + ) + assert await decide_export(request) is ExportActivationDecision.DENY + context.register_event_metadata_injector.assert_called_once() injector_name, injector = context.register_event_metadata_injector.call_args.args assert injector_name == "example_event_metadata_injector" diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py index 865590cda..1693218c5 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -27,6 +27,10 @@ ConfigDiagnostic, DataSchema, DiagnosticLevel, + ExportActivationDecision, + ExportActivationPolicyConfig, + ExportActivationTargetKind, + ExportTargetRegistration, Json, LlmOptimizationContribution, LlmRequestInterceptOutcome, @@ -50,6 +54,7 @@ from nemo_relay_plugin._api import ( # noqa: E402 ANNOTATED_LLM_REQUEST_SCHEMA, EVENT_SCHEMA, + EXPORT_ACTIVATION_REQUEST_SCHEMA, JSON_SCHEMA, LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, LLM_REQUEST_SCHEMA, @@ -535,6 +540,13 @@ async def subscriber(event: Json) -> None: async def event_metadata(event: Json) -> dict[str, Json]: return {"worker.event_name": event["name"]} + async def export_activation(request: Any) -> ExportActivationDecision: + return ( + ExportActivationDecision.ALLOW + if request.config.get("enabled") is True + else ExportActivationDecision.DENY + ) + async def mark_sanitize(event: Json, fields: Json) -> Json: return {**fields, "data": {"sanitized": f"mark:{event['name']}"}, "metadata": None} @@ -593,6 +605,18 @@ async def llm_stream_execution(name: str, request: Json, next_call: Any) -> Asyn ctx.register_subscriber("subscriber", subscriber) ctx.register_event_metadata_injector("event_metadata", event_metadata, priority=1) + ctx.register_export_activation_policy(export_activation) + ctx.register_export_target( + ExportTargetRegistration( + id="self_otel", + target_kind=ExportActivationTargetKind.OTLP_TRACE, + activation_policy=ExportActivationPolicyConfig( + provider=self.plugin_id, + config={"enabled": True}, + ), + ), + lambda: None, + ) ctx.register_mark_sanitize_guardrail("event_sanitize", mark_sanitize, priority=1) ctx.register_scope_sanitize_start_guardrail("event_sanitize", scope_start_sanitize, priority=2) ctx.register_scope_sanitize_end_guardrail("scope_end_sanitize", scope_end_sanitize, priority=3) @@ -637,6 +661,8 @@ def test_generated_proto_matches_worker_contract(): assert pb.HealthRequest.DESCRIPTOR.fields_by_name["auth_token"].number == 2 assert pb.SUBSCRIBER == 1 assert pb.EVENT_METADATA_INJECTOR == 2 + assert pb.EXPORT_ACTIVATION_POLICY == 3 + assert pb.InvokeRequest.DESCRIPTOR.fields_by_name["export_activation"].number == 13 assert pb.TOOL_SANITIZE_REQUEST_GUARDRAIL == 10 assert pb.LLM_STREAM_EXECUTION_INTERCEPT == 25 assert pb.MARK_SANITIZE_GUARDRAIL == 30 @@ -701,6 +727,8 @@ async def test_health_handshake_validate_register_and_all_surfaces(service: _Wor assert registrations == [ ("subscriber", pb.SUBSCRIBER, 0, False), ("event_metadata", pb.EVENT_METADATA_INJECTOR, 1, False), + ("export_activation_policy", pb.EXPORT_ACTIVATION_POLICY, 0, False), + ("self_otel", pb.EXPORT_TARGET, 0, False), ("event_sanitize", pb.MARK_SANITIZE_GUARDRAIL, 1, False), ("event_sanitize", pb.SCOPE_SANITIZE_START_GUARDRAIL, 2, False), ("scope_end_sanitize", pb.SCOPE_SANITIZE_END_GUARDRAIL, 3, False), @@ -1638,6 +1666,29 @@ async def test_unary_invoke_success_paths(service: _WorkerService, host_stub: Re assert event_metadata.WhichOneof("result") == "json" assert _envelope_value(event_metadata.json.value) == {"worker.event_name": "metadata-event"} + export_activation = await service.Invoke( + _invoke_request( + "export_activation_policy", + pb.EXPORT_ACTIVATION_POLICY, + export_activation=_json_envelope( + EXPORT_ACTIVATION_REQUEST_SCHEMA, + {"target_kind": "nemo_relay.otlp.trace", "config": {"enabled": True}}, + ), + ), + AbortContext(), + ) + assert _envelope_value(export_activation.json.value) == "allow" + + export_target = await service.Invoke( + _invoke_request( + "self_otel", + pb.EXPORT_TARGET, + export_target=_json_envelope("nemo.relay.ExportTargetActivation@1", None), + ), + AbortContext(), + ) + assert export_target.WhichOneof("result") == "empty" + tool_sanitize_request = await _invoke_json_async(service, "tool_sanitize", pb.TOOL_SANITIZE_REQUEST_GUARDRAIL) assert tool_sanitize_request["tag"] == "sanitize_lookup" tool_sanitize_response = await _invoke_json_async(service, "tool_sanitize", pb.TOOL_SANITIZE_RESPONSE_GUARDRAIL) @@ -3205,6 +3256,8 @@ def _all_expected_surfaces() -> list[int]: return [ pb.SUBSCRIBER, pb.EVENT_METADATA_INJECTOR, + pb.EXPORT_ACTIVATION_POLICY, + pb.EXPORT_TARGET, pb.MARK_SANITIZE_GUARDRAIL, pb.SCOPE_SANITIZE_START_GUARDRAIL, pb.SCOPE_SANITIZE_END_GUARDRAIL, diff --git a/python/tests/test_observability_plugin.py b/python/tests/test_observability_plugin.py index 07796c097..8780736ee 100644 --- a/python/tests/test_observability_plugin.py +++ b/python/tests/test_observability_plugin.py @@ -23,6 +23,7 @@ AtofFileSinkConfig, AtofStreamSinkConfig, ComponentSpec, + ExportActivationPolicyConfig, HttpStorageConfig, ObservabilityConfig, OpenTelemetryEndpointConfig, @@ -87,6 +88,42 @@ def wait_for_requests(self, expected: int, timeout: float = 5.0) -> list[tuple[d class TestObservabilityConfigHelpers: + def test_remote_targets_serialize_export_activation_policy(self): + policy = ExportActivationPolicyConfig( + provider="com.example.runtime-policy", + config={"allowed_countries": ["US", "CA"]}, + ) + expected = { + "provider": "com.example.runtime-policy", + "timeout_millis": 30000, + "config": {"allowed_countries": ["US", "CA"]}, + } + assert ( + OpenTelemetryEndpointConfig("full", "http://localhost:4318/v1/traces", activation_policy=policy).to_dict()[ + "activation_policy" + ] + == expected + ) + assert ( + OpenTelemetrySignalEndpointConfig("http://localhost:4318/v1/logs", activation_policy=policy).to_dict()[ + "activation_policy" + ] + == expected + ) + assert ( + AtofStreamSinkConfig("http://localhost:8080/events", activation_policy=policy).to_dict()[ + "activation_policy" + ] + == expected + ) + assert AtofFileSinkConfig(activation_policy=policy).to_dict()["activation_policy"] == expected + assert AtifConfig(local_activation_policy=policy).to_dict()["local_activation_policy"] == expected + assert ( + HttpStorageConfig("https://example.com/atif", activation_policy=policy).to_dict()["activation_policy"] + == expected + ) + assert S3StorageConfig("archive", activation_policy=policy).to_dict()["activation_policy"] == expected + def test_opentelemetry_endpoint_preserves_existing_positional_arguments(self): endpoint = OpenTelemetryEndpointConfig( "full", diff --git a/python/tests/test_plugin_export_activation.py b/python/tests/test_plugin_export_activation.py new file mode 100644 index 000000000..5ae58b3c8 --- /dev/null +++ b/python/tests/test_plugin_export_activation.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Ordinary Python plugin coverage for export activation hooks.""" + +from __future__ import annotations + +import pytest + +from nemo_relay import plugin + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("decision", "expected"), [("allow", 1), ("deny", 0)]) +async def test_plugin_can_gate_its_own_export_target(decision: str, expected: int) -> None: + kind = f"tests.python_export_activation_{decision}" + activated: list[str] = [] + + class SelfGatedPlugin: + def validate(self, plugin_config): + del plugin_config + return [] + + def register(self, plugin_config, context): + del plugin_config + + async def policy(request): + assert request == { + "target_kind": "tests.telemetry.otlp", + "config": {"country": "US"}, + } + return decision + + async def activate(): + activated.append("exporter") + + context.register_export_activation_policy(policy) + context.register_export_target( + { + "id": "self-otel", + "target_kind": "tests.telemetry.otlp", + "activation_policy": { + "provider": kind, + "timeout_millis": 30_000, + "config": {"country": "US"}, + }, + }, + activate, + ) + + plugin.register(kind, SelfGatedPlugin()) + try: + await plugin.initialize(plugin.PluginConfig(components=[plugin.ComponentSpec(kind=kind, config={})])) + assert len(activated) == expected + finally: + await plugin.clear_async() + plugin.deregister(kind)