From 445d3d654890ab66036031355bb4df85e3a538a4 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 20 Aug 2026 07:05:57 -0700 Subject: [PATCH 1/9] feat(observability): add dynamic export activation policies Signed-off-by: Bryan Bednarski --- crates/cli/src/server/mod.rs | 46 ++- crates/core/src/api/export_activation.rs | 70 ++++ crates/core/src/api/mod.rs | 2 + .../src/observability/plugin_component.rs | 273 +++++++++++++++- crates/core/src/plugin.rs | 31 ++ crates/core/src/plugin/dynamic.rs | 2 + crates/core/src/plugin/dynamic/host.rs | 34 +- crates/core/src/plugin/dynamic/native.rs | 43 +++ crates/core/src/plugin/dynamic/worker.rs | 51 +++ .../observability/plugin_component_tests.rs | 299 +++++++++++++++++- .../core/tests/unit/plugin_dynamic_tests.rs | 2 + crates/node/observability.d.ts | 15 + crates/node/observability.js | 25 ++ .../node/tests/observability_plugin_tests.mjs | 25 ++ .../tests/public_observability_api_fixture.ts | 15 + crates/plugin/src/async_sdk.rs | 27 ++ crates/plugin/src/lib.rs | 8 +- crates/plugin/tests/typed_callbacks.rs | 42 ++- crates/plugin/tests/unit/async_sdk_tests.rs | 4 + crates/types/src/plugin.rs | 43 +++ .../nemo/relay/worker/v1/plugin_worker.proto | 2 + crates/worker/README.md | 3 +- crates/worker/src/lib.rs | 69 +++- crates/worker/tests/worker_sdk_tests.rs | 52 ++- docs/build-plugins/dynamic-plugins/about.mdx | 8 + .../grpc-worker/python/about.mdx | 38 +++ .../rust-native-plugin-example.mdx | 36 +++ .../observability/configuration.mdx | 46 +++ examples/python-grpc-worker-plugin/README.md | 4 +- .../worker.py | 13 + .../relay-plugin.toml | 4 +- examples/rust-native-plugin/README.md | 2 + examples/rust-native-plugin/relay-plugin.toml | 2 +- examples/rust-native-plugin/src/lib.rs | 18 +- go/nemo_relay/observability_plugin.go | 153 +++++---- go/nemo_relay/observability_plugin_test.go | 33 +- justfile | 2 + python/nemo_relay/observability.py | 28 ++ python/nemo_relay/observability.pyi | 14 +- python/plugin/README.md | 2 + .../plugin/src/nemo_relay_plugin/__init__.py | 12 + python/plugin/src/nemo_relay_plugin/_api.py | 70 ++++ .../plugin/test_public_api_docstrings.py | 1 + .../plugin/test_python_worker_example.py | 21 +- python/tests/plugin/test_worker_sdk.py | 27 ++ python/tests/test_observability_plugin.py | 35 ++ 46 files changed, 1611 insertions(+), 141 deletions(-) create mode 100644 crates/core/src/api/export_activation.rs 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..8564e6663 --- /dev/null +++ b/crates/core/src/api/export_activation.rs @@ -0,0 +1,70 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Activation-time policy hooks for Relay-managed remote exporters. + +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::{Arc, LazyLock, RwLock}; + +use crate::error::{FlowError, Result}; +pub use nemo_relay_types::plugin::{ + ExportActivationDecision, ExportActivationRequest, ExportActivationTargetKind, +}; + +/// Asynchronous callback registered by one export-activation policy provider. +pub type ExportActivationPolicyFn = Arc< + dyn Fn( + ExportActivationRequest, + ) -> Pin> + Send>> + + Send + + Sync, +>; + +static EXPORT_ACTIVATION_POLICIES: LazyLock>> = + LazyLock::new(|| RwLock::new(HashMap::new())); + +pub(crate) fn register_export_activation_policy( + provider: &str, + callback: ExportActivationPolicyFn, +) -> Result<()> { + let mut policies = EXPORT_ACTIVATION_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(()) +} + +pub(crate) fn deregister_export_activation_policy(provider: &str) -> Result { + EXPORT_ACTIVATION_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_export_activation_policy( + provider: &str, + request: ExportActivationRequest, +) -> Result { + let callback = EXPORT_ACTIVATION_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 +} 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..701b2759b 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -36,6 +36,10 @@ use serde_json::{Map, Value as Json}; use uuid::Uuid; use crate::api::event::{Event, LogSeverity, ScopeCategory, ValidatedMetricMeasurement}; +use crate::api::export_activation::{ + ExportActivationDecision, ExportActivationRequest, ExportActivationTargetKind, + evaluate_export_activation_policy, +}; use crate::api::runtime::{EventSubscriberFn, current_scope_stack, global_context}; use crate::api::scope::ScopeType; use crate::api::subscriber::{ @@ -189,6 +193,20 @@ pub struct OpenTelemetrySectionConfig { pub metrics: Option, } +/// Activation-time policy attached to one remote exporter target. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct ExportActivationPolicyConfig { + /// Dynamic plugin identifier that owns the policy 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, +} + /// Signal-common OTLP destination fields used by logs and metrics. #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] @@ -223,6 +241,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 +384,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. @@ -436,6 +460,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. @@ -573,6 +600,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 +625,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 +800,7 @@ impl EditorConfig for OpenTelemetryEndpointConfig { &[], false, ), + otel_editor_field("activation_policy", EditorFieldKind::Json, &[], true), ], }; &SCHEMA @@ -797,6 +831,7 @@ impl EditorConfig for OpenTelemetrySignalEndpointConfig { &[], false, ), + otel_editor_field("activation_policy", EditorFieldKind::Json, &[], true), ], }; &SCHEMA @@ -870,6 +905,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 }, } } @@ -962,7 +998,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 +1100,7 @@ fn string_enum_schema( schema.into() } -fn register_observability( +async fn register_observability( config: ObservabilityConfig, ctx: &mut PluginRegistrationContext, ) -> PluginResult<()> { @@ -1087,13 +1123,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,10 +1157,31 @@ fn register_full_payload_policy( Ok(()) } -fn register_atof_exporter( - section: AtofSectionConfig, +async fn register_atof_exporter( + mut section: AtofSectionConfig, ctx: &mut PluginRegistrationContext, ) -> PluginResult<()> { + 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(_) => true, + AtofSinkSectionConfig::Stream(stream) => { + export_target_allowed( + stream.activation_policy.as_ref(), + ExportActivationTargetKind::AtofStream, + &format!("atof.sinks[{index}]"), + ) + .await + } + }; + if allowed { + allowed_sinks.push(sink); + } + } + section.sinks = allowed_sinks; + if section.sinks.is_empty() { + return Ok(()); + } let exporters = section .sinks .into_iter() @@ -1214,13 +1271,35 @@ 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<()> { validate_atif_filename_template(§ion.filename_template) .map_err(PluginError::InvalidConfig)?; + let had_remote_storage = !section.storage.is_empty(); + 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(policy, kind, &format!("atif.storage[{index}]")).await { + allowed_storage.push(entry); + } + } + section.storage = allowed_storage; + if had_remote_storage && section.storage.is_empty() { + return Ok(()); + } + let mut storage_vec = Vec::with_capacity(section.storage.len()); for (index, entry) in section.storage.iter().enumerate() { storage_vec.push(build_atif_storage(index, entry)?); @@ -1364,7 +1443,7 @@ fn build_atif_storage( )) } -fn register_opentelemetry( +async fn register_opentelemetry( section: OpenTelemetrySectionConfig, ctx: &mut PluginRegistrationContext, ) -> PluginResult<()> { @@ -1391,10 +1470,46 @@ 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( + 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( + 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( + 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 +1519,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 +1534,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 +1869,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(), }) } @@ -3263,6 +3382,7 @@ fn validate_opentelemetry_signal_fields( "service_version", "instrumentation_scope", "timeout_millis", + "activation_policy", ]; let section_fields = match signal { "logs" => &[ @@ -3344,6 +3464,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 +3687,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 +3870,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, @@ -4303,6 +4436,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( @@ -4539,6 +4678,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 +4725,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 +4757,40 @@ 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 !(1..=60_000).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 1 and 60000".to_string(), + ); + } +} + fn validate_atif_http_header( diagnostics: &mut Vec, policy: &ConfigPolicy, @@ -4797,10 +4982,64 @@ fn observability_registration_error(error: impl std::fmt::Display) -> PluginErro PluginError::RegistrationFailed(error.to_string()) } +async fn export_target_allowed( + policy: Option<&ExportActivationPolicyConfig>, + target_kind: ExportActivationTargetKind, + field: &str, +) -> bool { + let Some(policy) = policy else { + return true; + }; + let request = ExportActivationRequest { + target_kind, + config: policy.config.clone(), + }; + let outcome = tokio::time::timeout( + Duration::from_millis(policy.timeout_millis), + evaluate_export_activation_policy(&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 = export_activation_target_kind_name(target_kind), + field, + reason; + "Remote exporter target suppressed by activation policy" + ); + } + allowed +} + +const fn export_activation_target_kind_name(kind: ExportActivationTargetKind) -> &'static str { + match kind { + ExportActivationTargetKind::OtlpTrace => "otlp_trace", + ExportActivationTargetKind::OtlpLog => "otlp_log", + ExportActivationTargetKind::OtlpMetric => "otlp_metric", + ExportActivationTargetKind::AtofStream => "atof_stream", + ExportActivationTargetKind::AtifHttp => "atif_http", + ExportActivationTargetKind::AtifS3 => "atif_s3", + } +} + fn default_observability_config_version() -> u32 { 4 } +fn default_export_activation_timeout_millis() -> u64 { + 5_000 +} + fn default_atof_mode() -> String { "append".to_string() } diff --git a/crates/core/src/plugin.rs b/crates/core/src/plugin.rs index 4c37f2c16..1cac50add 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, deregister_export_activation_policy, + register_export_activation_policy, +}; use crate::api::registry::{ deregister_event_metadata_injector, deregister_llm_conditional_execution_guardrail, deregister_llm_execution_intercept, deregister_llm_request_intercept, @@ -433,6 +437,33 @@ impl PluginRegistrationContext { } } + /// Registers the single export-activation policy owned by `provider`. + pub fn register_export_activation_policy( + &mut self, + provider: &str, + callback: ExportActivationPolicyFn, + ) -> Result<()> { + register_export_activation_policy(provider, callback).map_err(|error| { + PluginError::RegistrationFailed(format!("export activation policy: {error}")) + })?; + + let provider = provider.to_string(); + self.registrations.push(PluginRegistration::new( + "export_activation_policy", + provider.clone(), + Box::new(move || { + deregister_export_activation_policy(&provider) + .map(|_| ()) + .map_err(|error| { + PluginError::RegistrationFailed(format!( + "export activation policy 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); diff --git a/crates/core/src/plugin/dynamic.rs b/crates/core/src/plugin/dynamic.rs index 6fbedc526..61e5fc799 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 Relay-managed remote exporters. + 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..8fef11171 100644 --- a/crates/core/src/plugin/dynamic/host.rs +++ b/crates/core/src/plugin/dynamic/host.rs @@ -23,7 +23,8 @@ use crate::plugin::{ }; use super::{ - DynamicPluginKind, DynamicPluginTeardownOutcome, NativePluginActivation, NativePluginLoadSpec, + DynamicPluginCapability, DynamicPluginKind, DynamicPluginManifest, + DynamicPluginTeardownOutcome, NativePluginActivation, NativePluginLoadSpec, load_native_plugins, }; @@ -177,15 +178,28 @@ 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 owner_id = claim.owner_id(); let initialization = tokio::spawn(initialize_plugins_exact_for_host( diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 4bb6e1b99..886cecb52 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -22,6 +22,9 @@ 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, +}; use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; use crate::api::runtime::{ EventMetadataInjectorFn, EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, @@ -3504,6 +3507,39 @@ 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_tool_execution( instance: Arc, cb: NemoRelayNativeAsyncMiddlewareCb, @@ -3813,6 +3849,13 @@ unsafe extern "C" fn native_plugin_context_register_async_middleware( priority, wrap_native_async_event_metadata_injector(instance, cb, user_data, free_fn), ), + NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy => { + let provider = instance.plugin_kind.clone(); + context.register_export_activation_policy( + &provider, + wrap_native_async_export_activation_policy(instance, cb, user_data, free_fn), + ) + } }; match registration { Ok(()) => NemoRelayStatus::Ok, diff --git a/crates/core/src/plugin/dynamic/worker.rs b/crates/core/src/plugin/dynamic/worker.rs index 4d276db5d..586e7901b 100644 --- a/crates/core/src/plugin/dynamic/worker.rs +++ b/crates/core/src/plugin/dynamic/worker.rs @@ -64,6 +64,9 @@ 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, +}; 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, @@ -96,6 +99,7 @@ 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 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); @@ -1071,6 +1075,8 @@ impl WorkerPluginInstance { ®istration.local_name, registration.priority, )?, + RegistrationSurface::ExportActivationPolicy => self + .install_export_activation_policy_registration(ctx, ®istration.local_name)?, RegistrationSurface::MarkSanitizeGuardrail | RegistrationSurface::ScopeSanitizeStartGuardrail | RegistrationSurface::ScopeSanitizeEndGuardrail => self @@ -1143,6 +1149,26 @@ 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 provider = self.plugin_kind.clone(); + 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(&provider, policy) + } + fn install_event_sanitize_registration( &self, ctx: &mut PluginRegistrationContext, @@ -1545,6 +1571,31 @@ 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_event_sanitize( &self, registration_name: &str, diff --git a/crates/core/tests/unit/observability/plugin_component_tests.rs b/crates/core/tests/unit/observability/plugin_component_tests.rs index 0446c8ee7..ec7b3169d 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,272 @@ 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( + "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::new(); + 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( + "test.slow-export-policy", + Arc::new(|_| { + Box::pin(async { + tokio::time::sleep(Duration::from_millis(25)).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, + config: Json::Null, + }); + let mut context = PluginRegistrationContext::new(); + 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( + 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::new(); + 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 duplicate_export_activation_provider_is_rejected_and_rollback_releases_it() { + 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("test.duplicate-policy", callback.clone()) + .unwrap(); + let mut duplicate = PluginRegistrationContext::new(); + assert!( + duplicate + .register_export_activation_policy("test.duplicate-policy", callback.clone()) + .is_err() + ); + + let mut first_registrations = first.into_registrations(); + rollback_registrations(&mut first_registrations); + duplicate + .register_export_activation_policy("test.duplicate-policy", callback) + .unwrap(); + let mut duplicate_registrations = duplicate.into_registrations(); + rollback_registrations(&mut duplicate_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(); + let section = AtofSectionConfig { + enabled: true, + sinks: vec![ + AtofSinkSectionConfig::File(AtofFileSinkSectionConfig { + output_directory: Some(temp_dir("allowed-local-atof")), + filename: Some("events.jsonl".into()), + mode: "overwrite".into(), + }), + AtofSinkSectionConfig::Stream(AtofStreamSinkSectionConfig { + name: Some("denied".into()), + url: "https://collector.example/events".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: 5_000, + config: Json::Null, + }), + }), + ], + }; + let mut context = PluginRegistrationContext::new(); + register_atof_exporter(section, &mut context).await.unwrap(); + let mut registrations = context.into_registrations(); + assert!(!registrations.is_empty()); + rollback_registrations(&mut registrations); +} + +#[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": 60001 + } + }] + } + }); + 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") + })); +} 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/node/observability.d.ts b/crates/node/observability.d.ts index e80f6b484..81ec4444a 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[]; @@ -29,6 +35,7 @@ export interface AtofStreamSinkConfig { timeout_millis?: number; field_name_policy?: 'preserve' | 'replace_dots' | string; name?: string; + activation_policy?: ExportActivationPolicyConfig; } /** @deprecated Use AtofStreamSinkConfig. */ @@ -44,6 +51,7 @@ export interface S3StorageConfig { region?: string; endpoint_url?: string; allow_http?: boolean; + activation_policy?: ExportActivationPolicyConfig; } export interface HttpStorageConfig { @@ -52,6 +60,7 @@ export interface HttpStorageConfig { headers?: Record; header_env?: Record; timeout_millis?: number; + activation_policy?: ExportActivationPolicyConfig; } export interface AtifConfig { @@ -85,6 +94,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 +108,7 @@ export interface OpenTelemetrySignalEndpointConfig { service_version?: string; instrumentation_scope?: string; timeout_millis?: number; + activation_policy?: ExportActivationPolicyConfig; } export interface OpenTelemetryLogSectionConfig { @@ -149,6 +160,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 remote-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..6268fcfee 100644 --- a/crates/node/observability.js +++ b/crates/node/observability.js @@ -47,6 +47,30 @@ function atifConfig(config = {}) { }; } +/** + * Create and validate one remote-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 ?? 5000; + if (!Number.isInteger(timeout) || timeout < 1 || timeout > 60000) { + throw new TypeError('Export activation policy timeout_millis must be between 1 and 60000'); + } + return { + timeout_millis: 5000, + config: null, + ...config, + }; +} + /** * Create one typed OpenTelemetry endpoint. * @@ -167,6 +191,7 @@ module.exports = { defaultConfig, atofConfig, atifConfig, + exportActivationPolicy, openTelemetryEndpoint, openTelemetrySignalEndpoint, openTelemetryLogConfig, diff --git a/crates/node/tests/observability_plugin_tests.mjs b/crates/node/tests/observability_plugin_tests.mjs index 46079a33d..42c3d604d 100644 --- a/crates/node/tests/observability_plugin_tests.mjs +++ b/crates/node/tests/observability_plugin_tests.mjs @@ -18,6 +18,31 @@ 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: 5000, + config: { allowed_countries: ['US', 'CA'] }, + }); + assert.equal( + observability.openTelemetryEndpoint({ + type: 'full', + endpoint: 'http://localhost:4318/v1/traces', + activation_policy: policy, + }).activation_policy, + policy, + ); + assert.throws(() => observability.exportActivationPolicy({ provider: ' ' }), /nonblank/); + assert.throws( + () => observability.exportActivationPolicy({ provider: 'test', timeout_millis: 60001 }), + /between 1 and 60000/, + ); + }); + it('builds defaults and plugin component shape', () => { assert.deepEqual(observability.defaultConfig(), { version: 4 }); assert.equal( 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..27febd116 100644 --- a/crates/plugin/src/async_sdk.rs +++ b/crates/plugin/src/async_sdk.rs @@ -1165,6 +1165,32 @@ impl PluginContext<'_> { ) } + /// Registers this plugin's activation-time policy for remote exporters. + 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_unary_adapter( + NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy, + "export_activation_policy", + 0, + false, + 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 an asynchronous mark-event sanitizer. pub fn register_mark_sanitize_guardrail( &mut self, @@ -1609,6 +1635,7 @@ fn registration_operation(kind: NemoRelayNativeAsyncMiddlewareKind) -> &'static NemoRelayNativeAsyncMiddlewareKind::ScopeSanitizeStart => "scope start sanitizer", NemoRelayNativeAsyncMiddlewareKind::ScopeSanitizeEnd => "scope end sanitizer", NemoRelayNativeAsyncMiddlewareKind::EventMetadataInjector => "Event metadata injector", + NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy => "export activation policy", } } diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 5028bc9ec..9334f4802 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, ExportActivationRequest, + ExportActivationTargetKind, +}; use serde::{Serialize, de::DeserializeOwned}; use serde_json::Map; @@ -828,6 +831,8 @@ pub enum NemoRelayNativeAsyncMiddlewareKind { ScopeSanitizeEnd = 13, /// Event metadata injector. EventMetadataInjector = 14, + /// Activation-time policy for Relay-managed remote exporters. + ExportActivationPolicy = 15, } impl TryFrom for NemoRelayNativeAsyncMiddlewareKind { @@ -850,6 +855,7 @@ impl TryFrom for NemoRelayNativeAsyncMiddlewareKind { 12 => Ok(Self::ScopeSanitizeStart), 13 => Ok(Self::ScopeSanitizeEnd), 14 => Ok(Self::EventMetadataInjector), + 15 => Ok(Self::ExportActivationPolicy), _ => Err(()), } } diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 6489a724f..b472ae8bd 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -22,15 +22,16 @@ use std::time::{Duration, Instant}; use futures::StreamExt; use nemo_relay_plugin::{ AnnotatedLlmRequest, BuiltinLlmCodec, CategoryProfile, ConfigDiagnostic, DataSchema, - DiagnosticLevel, Event, EventCategory, EventSanitizeFields, 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, + DiagnosticLevel, Event, EventCategory, EventSanitizeFields, ExportActivationDecision, + ExportActivationTargetKind, 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, NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV4, NemoRelayNativeLlmAsyncStream, @@ -68,12 +69,13 @@ fn async_abi_discriminants_reject_unknown_values() { Kind::ScopeSanitizeStart, Kind::ScopeSanitizeEnd, Kind::EventMetadataInjector, + Kind::ExportActivationPolicy, ]; for (discriminant, kind) in middleware_kinds.into_iter().enumerate() { assert_eq!(kind as u32, discriminant as u32); assert_eq!(Kind::try_from(discriminant as u32), Ok(kind)); } - assert!(NemoRelayNativeAsyncMiddlewareKind::try_from(15).is_err()); + assert!(NemoRelayNativeAsyncMiddlewareKind::try_from(16).is_err()); assert_eq!( NemoRelayNativeAsyncCallbackState::try_from(1), Ok(NemoRelayNativeAsyncCallbackState::Pending) @@ -3295,6 +3297,12 @@ 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(); let metadata = ASYNC_REGISTRATIONS .lock() @@ -3309,7 +3317,7 @@ fn typed_async_middleware_registers_and_round_trips_every_surface() { ) }) .collect::>(); - assert_eq!(metadata.len(), 14); + assert_eq!(metadata.len(), 15); assert_eq!(metadata[0].1, "mark-async"); assert_eq!(metadata[0].2, 1); assert_eq!( @@ -3353,6 +3361,18 @@ 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_async_registration(NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy); + let result = invoke_async_registration( + &host, + &policy_registration, + json!({"target_kind": "otlp_trace", "config": {"enabled": true}}), + None, + ) + .unwrap(); + assert_eq!(result, json!("allow")); + unsafe { policy_registration.free() }; + for kind in [ NemoRelayNativeAsyncMiddlewareKind::MarkSanitize, NemoRelayNativeAsyncMiddlewareKind::ScopeSanitizeStart, diff --git a/crates/plugin/tests/unit/async_sdk_tests.rs b/crates/plugin/tests/unit/async_sdk_tests.rs index 4d389674e..234c67319 100644 --- a/crates/plugin/tests/unit/async_sdk_tests.rs +++ b/crates/plugin/tests/unit/async_sdk_tests.rs @@ -125,6 +125,10 @@ fn async_sdk_status_helpers_keep_operation_labels_stable() { NemoRelayNativeAsyncMiddlewareKind::EventMetadataInjector, "Event metadata injector", ), + ( + NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy, + "export activation policy", + ), ]; for (kind, expected) in labels { assert_eq!(registration_operation(kind), expected); diff --git a/crates/types/src/plugin.rs b/crates/types/src/plugin.rs index eeb3be024..6c66e73ec 100644 --- a/crates/types/src/plugin.rs +++ b/crates/types/src/plugin.rs @@ -5,6 +5,49 @@ use serde::{Deserialize, Serialize}; +use crate::Json; + +/// Remote exporter target presented to an activation policy. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum ExportActivationTargetKind { + /// OpenTelemetry trace destination. + OtlpTrace, + /// OpenTelemetry log destination. + OtlpLog, + /// OpenTelemetry metric destination. + OtlpMetric, + /// Remote ATOF stream sink. + AtofStream, + /// Remote ATIF HTTP storage. + AtifHttp, + /// Remote ATIF S3-compatible storage. + AtifS3, +} + +/// 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 remote exporter 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 remote exporter target. + Allow, + /// Suppress the remote exporter target for this activation. + Deny, +} + /// Diagnostic severity returned by plugin validation. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] 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..23146dbf3 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,7 @@ enum RegistrationSurface { REGISTRATION_SURFACE_UNSPECIFIED = 0; SUBSCRIBER = 1; EVENT_METADATA_INJECTOR = 2; + EXPORT_ACTIVATION_POLICY = 3; TOOL_SANITIZE_REQUEST_GUARDRAIL = 10; TOOL_SANITIZE_RESPONSE_GUARDRAIL = 11; TOOL_CONDITIONAL_EXECUTION_GUARDRAIL = 12; @@ -165,6 +166,7 @@ message InvokeRequest { JsonEnvelope event = 10; ToolInvocation tool = 11; LlmInvocation llm = 12; + JsonEnvelope export_activation = 13; } } diff --git a/crates/worker/README.md b/crates/worker/README.md index 5b29a2daa..496662d41 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 Relay-managed remote 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..0e7ae87ac 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, ExportActivationRequest, + ExportActivationTargetKind, +}; 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,7 @@ 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 DATA_SCHEMA_SCHEMA: &str = "nemo.relay.DataSchema@1"; const LLM_REQUEST_SCHEMA: &str = "nemo.relay.LlmRequest@1"; @@ -184,6 +188,8 @@ type EventSanitizeFn = Arc BoxFutureResult + Send + Sync>; type EventMetadataInjectorFn = Arc) -> BoxFutureResult> + Send + Sync>; +type ExportActivationPolicyFn = + Arc BoxFutureResult + Send + Sync>; type ToolSanitizeFn = Arc BoxFutureResult + Send + Sync>; type ToolConditionalFn = Arc BoxFutureResult> + Send + Sync>; type ToolRequestFn = Arc BoxFutureResult + Send + Sync>; @@ -309,6 +315,7 @@ struct WorkerHandlers { registrations: Vec, subscribers: HashMap, event_metadata_injectors: HashMap, + export_activation_policies: HashMap, mark_sanitizers: HashMap, scope_start_sanitizers: HashMap, scope_end_sanitizers: HashMap, @@ -386,6 +393,20 @@ impl PluginContext { ); } + /// Registers this plugin's activation-time policy for remote 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))), + ); + } + fn register_event_sanitizer( &mut self, name: &str, @@ -1819,6 +1840,10 @@ impl WorkerService { self.invoke_event_metadata_injector_response(request, &scope) .await } + RegistrationSurface::ExportActivationPolicy => { + self.invoke_export_activation_policy_response(request, &scope) + .await + } RegistrationSurface::MarkSanitizeGuardrail | RegistrationSurface::ScopeSanitizeStartGuardrail | RegistrationSurface::ScopeSanitizeEndGuardrail => { @@ -1871,6 +1896,33 @@ 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_event_sanitize_response( &self, request: InvokeRequest, @@ -2116,6 +2168,20 @@ 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 event_sanitizer(&self, surface: RegistrationSurface, name: &str) -> Result { let handlers = self .handlers @@ -2743,6 +2809,7 @@ fn all_surfaces() -> Vec { vec![ RegistrationSurface::Subscriber, RegistrationSurface::EventMetadataInjector, + RegistrationSurface::ExportActivationPolicy, 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..fb1a43833 100644 --- a/crates/worker/tests/worker_sdk_tests.rs +++ b/crates/worker/tests/worker_sdk_tests.rs @@ -20,11 +20,11 @@ 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, 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 +127,11 @@ 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)) + ); let bad_health = client .health(Request::new(HealthRequest { @@ -213,7 +218,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(), 23); for local_name in [ "llm-sanitize-request", "llm-sanitize-response", @@ -545,6 +550,9 @@ 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 mark_fields = invoke_json( &mut client, event_invoke_surface("event-sanitize", RegistrationSurface::MarkSanitizeGuardrail), @@ -1833,6 +1841,17 @@ 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_mark_sanitize_guardrail("event-sanitize", 1, |event, mut fields| { let event_name = event.name().to_owned(); async move { @@ -2688,6 +2707,27 @@ 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": "otlp_trace", "config": {"enabled": true}}), + ) + .expect("encode export activation request"), + ), + ), + } +} + 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..851d72e7b 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,44 @@ 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, + ExportActivationRequest, + ExportActivationTargetKind, +) + + +async def decide_export(request: ExportActivationRequest) -> ExportActivationDecision: + config = request.config if isinstance(request.config, dict) else {} + if request.target_kind in { + ExportActivationTargetKind.OTLP_TRACE, + ExportActivationTargetKind.OTLP_LOG, + ExportActivationTargetKind.OTLP_METRIC, + }: + return ( + ExportActivationDecision.ALLOW + if config.get("telemetry_allowed") is True + else ExportActivationDecision.DENY + ) + return ExportActivationDecision.ALLOW + + +ctx.register_export_activation_policy(decide_export) +``` + +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. + ## 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..114481b29 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,39 @@ 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, ExportActivationTargetKind, Json, PluginContext, Result, +}; + +fn register_policy(context: &mut PluginContext<'_>) -> Result<()> { + context.register_export_activation_policy(|request| async move { + let allowed = match request.target_kind { + ExportActivationTargetKind::OtlpTrace + | ExportActivationTargetKind::OtlpLog + | ExportActivationTargetKind::OtlpMetric => { + request.config.get("telemetry_allowed") == Some(&Json::Bool(true)) + } + _ => true, + }; + Ok(if allowed { + ExportActivationDecision::Allow + } else { + ExportActivationDecision::Deny + }) + }) +} +``` + +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. diff --git a/docs/configure-plugins/observability/configuration.mdx b/docs/configure-plugins/observability/configuration.mdx index 3face373c..dc3ea651f 100644 --- a/docs/configure-plugins/observability/configuration.mdx +++ b/docs/configure-plugins/observability/configuration.mdx @@ -89,6 +89,43 @@ 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 remote targets accept an optional `activation_policy`. The +`provider` is the dynamic-plugin manifest ID of a native or worker plugin that +declares the `export_activation_policy` capability and registers one callback: + +```toml +[components.config.opentelemetry.endpoints.activation_policy] +provider = "com.example.runtime-policy" +timeout_millis = 5000 +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. + +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 5000 milliseconds and must be between 1 and 60000. +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 +200,7 @@ from nemo_relay.observability import ( ObservabilityConfig, OpenTelemetryLogSectionConfig, OpenTelemetryMetricSectionConfig, + ExportActivationPolicyConfig, OpenTelemetrySectionConfig, OpenTelemetryEndpointConfig, ) @@ -178,6 +216,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 +245,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/examples/python-grpc-worker-plugin/README.md b/examples/python-grpc-worker-plugin/README.md index b24a3cb6c..d47f0895a 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 a remote 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..eb777cc4e 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 a remote 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/observability_plugin.go b/go/nemo_relay/observability_plugin.go index 987d71cfd..ef3e79029 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 dynamic policy provider to a remote exporter 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. @@ -118,13 +127,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() { @@ -164,14 +174,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 +192,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 +213,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 +233,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 +261,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: 5000} +} + // 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..f8ed333f6 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,13 +43,14 @@ 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" { @@ -63,6 +67,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 +82,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 +96,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 @@ -148,6 +156,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 +172,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 +199,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(5000) { + t.Fatalf("unexpected export activation policy: %#v", value) + } } func TestObservabilitySignalEndpointOmittedVersusExplicitEmpty(t *testing.T) { 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/observability.py b/python/nemo_relay/observability.py index 113ef1575..e4c32fba4 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 remote exporter target.""" + + provider: str + timeout_millis: int = 5000 + 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, } ) @@ -139,6 +158,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 +173,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 +186,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 +197,7 @@ def to_dict(self) -> JsonObject: "headers": self.headers, "header_env": self.header_env, "timeout_millis": self.timeout_millis, + "activation_policy": self.activation_policy, } ) @@ -233,6 +256,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 +280,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 +299,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 +315,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 +441,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..9ad4cf909 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) @@ -53,6 +61,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 +70,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) @@ -96,6 +106,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 +121,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/plugin/README.md b/python/plugin/README.md index 04e001eaf..e9deb1a88 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 remote 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..02fb5c230 100644 --- a/python/plugin/src/nemo_relay_plugin/__init__.py +++ b/python/plugin/src/nemo_relay_plugin/__init__.py @@ -45,6 +45,9 @@ ToolExecutionInterceptOutcome: Canonical tool execution-intercept result. DiagnosticLevel: Severity of a configuration diagnostic. ConfigDiagnostic: Structured configuration warning or error. + ExportActivationTargetKind: Remote exporter kind presented to a policy. + ExportActivationRequest: Secret-free activation request for one remote target. + ExportActivationDecision: Allow-or-deny result returned by an activation policy. RuntimeDiagnostic: One aggregated host runtime diagnostic entry. RuntimeDiagnostics: Immutable host runtime diagnostics snapshot. ScopeType: Semantic category for a Relay execution scope. @@ -53,6 +56,7 @@ Public callback aliases: SubscriberCallback: Event subscriber callback. EventMetadataInjectorCallback: Event metadata injector callback. + ExportActivationPolicyCallback: Activation-time remote-export policy callback. EventSanitizeCallback: Mark or scope event sanitizer callback. ToolSanitizeCallback: Tool request or response sanitizer callback. ToolConditionalCallback: Tool execution guardrail callback. @@ -86,6 +90,10 @@ EventMetadataInjectorCallback, EventSanitizeCallback, EventSanitizeFields, + ExportActivationDecision, + ExportActivationPolicyCallback, + ExportActivationRequest, + ExportActivationTargetKind, Json, LlmCodecIdentity, LlmConditionalCallback, @@ -141,6 +149,10 @@ "EventMetadataInjectorCallback", "EventSanitizeCallback", "EventSanitizeFields", + "ExportActivationDecision", + "ExportActivationPolicyCallback", + "ExportActivationRequest", + "ExportActivationTargetKind", "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..b9845a318 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: Remote exporter kind presented to a policy. + ExportActivationRequest: Secret-free activation request for one remote 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 remote-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,38 @@ 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, Enum): + """Remote exporter target presented to an activation policy.""" + + OTLP_TRACE = "otlp_trace" + OTLP_LOG = "otlp_log" + OTLP_METRIC = "otlp_metric" + ATOF_STREAM = "atof_stream" + ATIF_HTTP = "atif_http" + ATIF_S3 = "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 + + +ExportActivationPolicyCallback: TypeAlias = Callable[ + [ExportActivationRequest], + ExportActivationDecision | Awaitable[ExportActivationDecision], +] EventSanitizeCallback: TypeAlias = Callable[ [Event, EventSanitizeFields], EventSanitizeFields | Awaitable[EventSanitizeFields], @@ -971,6 +1008,7 @@ class _Handlers: registrations: list[Any] subscribers: dict[str, SubscriberCallback] event_metadata_injectors: dict[str, EventMetadataInjectorCallback] + export_activation_policies: dict[str, ExportActivationPolicyCallback] mark_sanitizers: dict[str, EventSanitizeCallback] scope_start_sanitizers: dict[str, EventSanitizeCallback] scope_end_sanitizers: dict[str, EventSanitizeCallback] @@ -992,6 +1030,7 @@ def empty(cls) -> _Handlers: registrations=[], subscribers={}, event_metadata_injectors={}, + export_activation_policies={}, mark_sanitizers={}, scope_start_sanitizers={}, scope_end_sanitizers={}, @@ -1097,6 +1136,12 @@ 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 remote exporters.""" + name = "export_activation_policy" + self._push_registration(name, pb.EXPORT_ACTIVATION_POLICY, 0, False) + self._handlers.export_activation_policies[name] = callback + def _register_event_sanitizer( self, name: str, @@ -2230,6 +2275,30 @@ 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 in PluginContext._EVENT_SANITIZER_HANDLER_ATTRIBUTES: event = _decode_required_envelope(request.event, "event", EVENT_SCHEMA) fields: EventSanitizeFields = { @@ -2414,6 +2483,7 @@ def _all_surfaces() -> list[int]: return [ pb.SUBSCRIBER, pb.EVENT_METADATA_INJECTOR, + pb.EXPORT_ACTIVATION_POLICY, 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..1ea3ece57 100644 --- a/python/tests/plugin/test_public_api_docstrings.py +++ b/python/tests/plugin/test_public_api_docstrings.py @@ -23,6 +23,7 @@ "Event", "EventMetadataInjectorCallback", "EventSanitizeCallback", + "ExportActivationPolicyCallback", "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..e46d91da3 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -27,6 +27,7 @@ ConfigDiagnostic, DataSchema, DiagnosticLevel, + ExportActivationDecision, Json, LlmOptimizationContribution, LlmRequestInterceptOutcome, @@ -50,6 +51,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 +537,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 +602,7 @@ 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_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 +647,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 +713,7 @@ 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), ("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 +1651,19 @@ 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": "otlp_trace", "config": {"enabled": True}}, + ), + ), + AbortContext(), + ) + assert _envelope_value(export_activation.json.value) == "allow" + 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 +3231,7 @@ def _all_expected_surfaces() -> list[int]: return [ pb.SUBSCRIBER, pb.EVENT_METADATA_INJECTOR, + pb.EXPORT_ACTIVATION_POLICY, 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..270e387df 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,40 @@ 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": 5000, + "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 ( + 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", From 9c92d477d57ea82e7713e14616fd6051f24fdea8 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 20 Aug 2026 08:43:21 -0700 Subject: [PATCH 2/9] fix(exporters): validate worker policy surface Signed-off-by: Bryan Bednarski --- crates/core/src/plugin/dynamic/worker.rs | 31 ++++++++++++++++++- .../core/tests/unit/dynamic_worker_tests.rs | 25 +++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/crates/core/src/plugin/dynamic/worker.rs b/crates/core/src/plugin/dynamic/worker.rs index 586e7901b..661d5386e 100644 --- a/crates/core/src/plugin/dynamic/worker.rs +++ b/crates/core/src/plugin/dynamic/worker.rs @@ -615,11 +615,16 @@ 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_policy( + &spec.plugin_id, + &handshake.supported_surfaces, + &mut register.registrations, + ); register.registrations }; if registrations.iter().any(|registration| { @@ -3391,6 +3396,30 @@ fn validate_registration_plan( Ok(()) } +fn filter_unadvertised_export_activation_policy( + plugin_id: &str, + supported_surfaces: &[i32], + registrations: &mut Vec, +) { + let policy_surface = RegistrationSurface::ExportActivationPolicy as i32; + if supported_surfaces.contains(&policy_surface) { + return; + } + registrations.retain(|registration| { + if registration.surface != policy_surface { + return true; + } + log::warn!( + target: "nemo_relay.worker", + event = "worker_registration_surface_unadvertised", + plugin_id, + surface = "export_activation_policy"; + "Worker export activation policy 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..06815fe13 100644 --- a/crates/core/tests/unit/dynamic_worker_tests.rs +++ b/crates/core/tests/unit/dynamic_worker_tests.rs @@ -561,6 +561,31 @@ 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_policy_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_policy( + "fixture_worker", + &[RegistrationSurface::Subscriber as i32], + &mut registrations, + ); + assert_eq!(registrations, vec![subscriber]); + + let mut advertised = vec![policy.clone()]; + filter_unadvertised_export_activation_policy( + "fixture_worker", + &[RegistrationSurface::ExportActivationPolicy as i32], + &mut advertised, + ); + assert_eq!(advertised, vec![policy]); +} + #[test] fn relay_compatibility_and_blocking_helpers_cover_local_edges() { enable_operational_logs(); From d6ecf39f7d848f32a1fca9a2f3c0ce7975466eb5 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 20 Aug 2026 09:46:39 -0700 Subject: [PATCH 3/9] refactor(plugins): scope export policies to activation Signed-off-by: Bryan Bednarski --- crates/core/src/api/export_activation.rs | 88 +++++----- .../src/observability/plugin_component.rs | 23 ++- crates/core/src/plugin.rs | 163 +++++++++++++++--- crates/core/src/plugin/dynamic/host.rs | 6 + .../observability/plugin_component_tests.rs | 23 ++- .../tests/unit/plugin_dynamic_host_tests.rs | 1 + crates/core/tests/unit/plugin_tests.rs | 7 +- .../observability/configuration.mdx | 7 + 8 files changed, 240 insertions(+), 78 deletions(-) diff --git a/crates/core/src/api/export_activation.rs b/crates/core/src/api/export_activation.rs index 8564e6663..a5bc9b35b 100644 --- a/crates/core/src/api/export_activation.rs +++ b/crates/core/src/api/export_activation.rs @@ -6,7 +6,7 @@ use std::collections::HashMap; use std::future::Future; use std::pin::Pin; -use std::sync::{Arc, LazyLock, RwLock}; +use std::sync::{Arc, RwLock}; use crate::error::{FlowError, Result}; pub use nemo_relay_types::plugin::{ @@ -22,49 +22,57 @@ pub type ExportActivationPolicyFn = Arc< + Sync, >; -static EXPORT_ACTIVATION_POLICIES: LazyLock>> = - LazyLock::new(|| RwLock::new(HashMap::new())); - -pub(crate) fn register_export_activation_policy( - provider: &str, - callback: ExportActivationPolicyFn, -) -> Result<()> { - let mut policies = EXPORT_ACTIVATION_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(()) +/// Export-activation policy callbacks owned by one plugin activation. +#[derive(Default)] +pub(crate) struct ExportActivationPolicyRegistry { + policies: RwLock>, } -pub(crate) fn deregister_export_activation_policy(provider: &str) -> Result { - EXPORT_ACTIVATION_POLICIES - .write() - .map(|mut policies| policies.remove(provider).is_some()) - .map_err(|error| { +impl ExportActivationPolicyRegistry { + pub(crate) 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(()) + } -pub(crate) async fn evaluate_export_activation_policy( - provider: &str, - request: ExportActivationRequest, -) -> Result { - let callback = EXPORT_ACTIVATION_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 + pub(crate) 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 + } } diff --git a/crates/core/src/observability/plugin_component.rs b/crates/core/src/observability/plugin_component.rs index 701b2759b..33bc7e4fe 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -37,8 +37,8 @@ use uuid::Uuid; use crate::api::event::{Event, LogSeverity, ScopeCategory, ValidatedMetricMeasurement}; use crate::api::export_activation::{ - ExportActivationDecision, ExportActivationRequest, ExportActivationTargetKind, - evaluate_export_activation_policy, + ExportActivationDecision, ExportActivationPolicyRegistry, ExportActivationRequest, + ExportActivationTargetKind, }; use crate::api::runtime::{EventSubscriberFn, current_scope_stack, global_context}; use crate::api::scope::ScopeType; @@ -1161,12 +1161,14 @@ async fn register_atof_exporter( mut section: AtofSectionConfig, ctx: &mut PluginRegistrationContext, ) -> PluginResult<()> { + 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(_) => true, AtofSinkSectionConfig::Stream(stream) => { export_target_allowed( + &export_activation_policies, stream.activation_policy.as_ref(), ExportActivationTargetKind::AtofStream, &format!("atof.sinks[{index}]"), @@ -1275,6 +1277,7 @@ 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)?; @@ -1291,7 +1294,14 @@ async fn register_atif_dispatcher( ExportActivationTargetKind::AtifS3, ), }; - if export_target_allowed(policy, kind, &format!("atif.storage[{index}]")).await { + if export_target_allowed( + &export_activation_policies, + policy, + kind, + &format!("atif.storage[{index}]"), + ) + .await + { allowed_storage.push(entry); } } @@ -1447,6 +1457,7 @@ async fn register_opentelemetry( section: OpenTelemetrySectionConfig, ctx: &mut PluginRegistrationContext, ) -> PluginResult<()> { + let export_activation_policies = ctx.export_activation_policies(); let OpenTelemetrySectionConfig { endpoints, logs, @@ -1473,6 +1484,7 @@ async fn register_opentelemetry( 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}]"), @@ -1485,6 +1497,7 @@ async fn register_opentelemetry( 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}]"), @@ -1497,6 +1510,7 @@ async fn register_opentelemetry( 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}]"), @@ -4983,6 +4997,7 @@ fn observability_registration_error(error: impl std::fmt::Display) -> PluginErro } async fn export_target_allowed( + export_activation_policies: &ExportActivationPolicyRegistry, policy: Option<&ExportActivationPolicyConfig>, target_kind: ExportActivationTargetKind, field: &str, @@ -4996,7 +5011,7 @@ async fn export_target_allowed( }; let outcome = tokio::time::timeout( Duration::from_millis(policy.timeout_millis), - evaluate_export_activation_policy(&policy.provider, request), + export_activation_policies.evaluate(&policy.provider, request), ) .await; let (allowed, reason) = match outcome { diff --git a/crates/core/src/plugin.rs b/crates/core/src/plugin.rs index 1cac50add..b0cba4032 100644 --- a/crates/core/src/plugin.rs +++ b/crates/core/src/plugin.rs @@ -22,10 +22,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value as Json}; use thiserror::Error; -use crate::api::export_activation::{ - ExportActivationPolicyFn, deregister_export_activation_policy, - register_export_activation_policy, -}; +use crate::api::export_activation::{ExportActivationPolicyFn, ExportActivationPolicyRegistry}; use crate::api::registry::{ deregister_event_metadata_injector, deregister_llm_conditional_execution_guardrail, deregister_llm_execution_intercept, deregister_llm_request_intercept, @@ -405,10 +402,20 @@ 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, + export_activation_policies: Arc, +} + +impl Default for PluginRegistrationContext { + fn default() -> Self { + Self { + registrations: Vec::new(), + namespace: None, + export_activation_policies: Arc::new(ExportActivationPolicyRegistry::default()), + } + } } impl PluginRegistrationContext { @@ -422,9 +429,36 @@ impl PluginRegistrationContext { Self { registrations: vec![], namespace: Some(namespace.into()), + export_activation_policies: Arc::new(ExportActivationPolicyRegistry::default()), + } + } + + pub(crate) fn with_namespace_and_export_activation_policies( + namespace: impl Into, + export_activation_policies: Arc, + ) -> Self { + Self { + registrations: Vec::new(), + namespace: Some(namespace.into()), + export_activation_policies, + } + } + + #[cfg(test)] + pub(crate) fn with_export_activation_policies( + export_activation_policies: Arc, + ) -> Self { + Self { + registrations: Vec::new(), + namespace: None, + export_activation_policies, } } + pub(crate) fn export_activation_policies(&self) -> Arc { + Arc::clone(&self.export_activation_policies) + } + /// Returns the runtime-qualified name for a plugin-local registration. /// /// Plugin handlers should pass stable component-local names such as @@ -443,16 +477,20 @@ impl PluginRegistrationContext { provider: &str, callback: ExportActivationPolicyFn, ) -> Result<()> { - register_export_activation_policy(provider, callback).map_err(|error| { - PluginError::RegistrationFailed(format!("export activation policy: {error}")) - })?; + 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 || { - deregister_export_activation_policy(&provider) + export_activation_policies + .deregister(&provider) .map(|_| ()) .map_err(|error| { PluginError::RegistrationFailed(format!( @@ -1517,10 +1555,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| { @@ -1638,15 +1678,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 @@ -1685,6 +1733,7 @@ async fn initialize_plugins_exact_inner( previous_state, rollback_failures, enabled_component_count, + export_activation_policies, ) .await } @@ -1694,6 +1743,7 @@ async fn initialize_plugins_exact_inner( report, rollback_failures, enabled_component_count, + export_activation_policies, ) .await } @@ -1705,10 +1755,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", @@ -1724,6 +1784,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); @@ -1741,6 +1802,7 @@ async fn replace_plugin_configuration( previous_state, rollback_failures, enabled_component_count, + export_activation_policies, ) .await } @@ -1756,6 +1818,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(()) } @@ -1791,12 +1854,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", @@ -1816,9 +1889,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 { @@ -1828,6 +1903,7 @@ async fn restore_previous_plugin_configuration( previous_state.report, previous_state.runtime_diagnostics, registrations, + export_activation_policies, )?; log::warn!( target: "nemo_relay.plugin", @@ -1854,14 +1930,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 @@ -2637,11 +2716,13 @@ 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); @@ -2670,8 +2751,11 @@ async fn initialize_plugin_components( totals.get(component.kind.as_str()).copied().unwrap_or(1), ); - let mut pending = - PendingPluginRegistrationContext::new(namespace, rollback_failures.clone()); + let mut pending = PendingPluginRegistrationContext::new( + namespace, + rollback_failures.clone(), + Arc::clone(&export_activation_policies), + ); plugin .register(&component.config, &mut pending.context) .await?; @@ -2718,9 +2802,16 @@ struct PendingPluginRegistrationContext { } impl PendingPluginRegistrationContext { - fn new(namespace: String, rollback_failures: Option>>>) -> Self { + fn new( + namespace: String, + rollback_failures: Option>>>, + export_activation_policies: Arc, + ) -> Self { Self { - context: PluginRegistrationContext::with_namespace(namespace), + context: PluginRegistrationContext::with_namespace_and_export_activation_policies( + namespace, + export_activation_policies, + ), rollback_failures, } } @@ -2753,16 +2844,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, ) } @@ -2771,6 +2878,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}")) @@ -2780,6 +2888,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/host.rs b/crates/core/src/plugin/dynamic/host.rs index 8fef11171..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, @@ -57,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, @@ -201,12 +203,14 @@ impl PluginHostActivation { 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| { @@ -261,6 +265,7 @@ impl PluginHostActivation { Ok(( Self { active: true, + export_activation_policies: Some(export_activation_policies), native, #[cfg(feature = "worker-grpc")] worker, @@ -309,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/tests/unit/observability/plugin_component_tests.rs b/crates/core/tests/unit/observability/plugin_component_tests.rs index ec7b3169d..50b9c24d3 100644 --- a/crates/core/tests/unit/observability/plugin_component_tests.rs +++ b/crates/core/tests/unit/observability/plugin_component_tests.rs @@ -5030,7 +5030,9 @@ async fn export_activation_policy_is_evaluated_once_before_otel_construction() { timeout_millis: 5_000, config: json!({"region": "local"}), }); - let mut observability_context = PluginRegistrationContext::new(); + let mut observability_context = PluginRegistrationContext::with_export_activation_policies( + provider_context.export_activation_policies(), + ); register_opentelemetry( OpenTelemetrySectionConfig { enabled: true, @@ -5092,7 +5094,9 @@ async fn unavailable_and_timed_out_export_policies_fail_closed() { timeout_millis: 1, config: Json::Null, }); - let mut context = PluginRegistrationContext::new(); + let mut context = PluginRegistrationContext::with_export_activation_policies( + provider_context.export_activation_policies(), + ); register_opentelemetry( OpenTelemetrySectionConfig { enabled: true, @@ -5141,7 +5145,9 @@ async fn denied_and_failed_export_policies_fail_closed() { timeout_millis: 5_000, config: Json::Null, }); - let mut context = PluginRegistrationContext::new(); + let mut context = PluginRegistrationContext::with_export_activation_policies( + provider_context.export_activation_policies(), + ); register_opentelemetry( OpenTelemetrySectionConfig { enabled: true, @@ -5160,19 +5166,24 @@ async fn denied_and_failed_export_policies_fail_closed() { } #[test] -fn duplicate_export_activation_provider_is_rejected_and_rollback_releases_it() { +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("test.duplicate-policy", callback.clone()) .unwrap(); - let mut duplicate = PluginRegistrationContext::new(); + let registry = first.export_activation_policies(); + let mut duplicate = PluginRegistrationContext::with_export_activation_policies(registry); assert!( duplicate .register_export_activation_policy("test.duplicate-policy", callback.clone()) .is_err() ); + let mut independent_activation = PluginRegistrationContext::new(); + independent_activation + .register_export_activation_policy("test.duplicate-policy", callback.clone()) + .unwrap(); let mut first_registrations = first.into_registrations(); rollback_registrations(&mut first_registrations); @@ -5181,6 +5192,8 @@ fn duplicate_export_activation_provider_is_rejected_and_rollback_releases_it() { .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] 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_tests.rs b/crates/core/tests/unit/plugin_tests.rs index 7816bfbb5..59b12c916 100644 --- a/crates/core/tests/unit/plugin_tests.rs +++ b/crates/core/tests/unit/plugin_tests.rs @@ -1735,8 +1735,11 @@ 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(), + Some(Arc::clone(&failures)), + Arc::new(ExportActivationPolicyRegistry::default()), + ); pending .context .add_registration(PluginRegistration::new_with_outcome( diff --git a/docs/configure-plugins/observability/configuration.mdx b/docs/configure-plugins/observability/configuration.mdx index dc3ea651f..430780039 100644 --- a/docs/configure-plugins/observability/configuration.mdx +++ b/docs/configure-plugins/observability/configuration.mdx @@ -108,6 +108,13 @@ that exporter. The callback receives only the target kind and the opaque 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 From cca388bab1263cd739da04a924262684015680f7 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 20 Aug 2026 09:49:52 -0700 Subject: [PATCH 4/9] fix(exporters): bound activation policy timeouts Signed-off-by: Bryan Bednarski --- .../src/observability/plugin_component.rs | 21 +++++++-- .../observability/plugin_component_tests.rs | 47 +++++++++++++++---- crates/node/observability.js | 8 ++-- .../node/tests/observability_plugin_tests.mjs | 10 ++-- .../observability/configuration.mdx | 7 ++- go/nemo_relay/observability_plugin.go | 2 +- go/nemo_relay/observability_plugin_test.go | 2 +- python/nemo_relay/observability.py | 2 +- python/tests/test_observability_plugin.py | 2 +- 9 files changed, 75 insertions(+), 26 deletions(-) diff --git a/crates/core/src/observability/plugin_component.rs b/crates/core/src/observability/plugin_component.rs index 33bc7e4fe..862bf6433 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -207,6 +207,10 @@ pub struct ExportActivationPolicyConfig { pub config: Json, } +const MIN_EXPORT_ACTIVATION_TIMEOUT_MILLIS: u64 = 1_000; +const DEFAULT_EXPORT_ACTIVATION_TIMEOUT_MILLIS: u64 = 30_000; +const MAX_EXPORT_ACTIVATION_TIMEOUT_MILLIS: u64 = 300_000; + /// Signal-common OTLP destination fields used by logs and metrics. #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] @@ -4793,14 +4797,16 @@ fn validate_export_activation_policy( .to_string(), ); } - if !(1..=60_000).contains(&activation_policy.timeout_millis) { + 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 1 and 60000".to_string(), + "export activation policy timeout_millis must be between 1000 and 300000".to_string(), ); } } @@ -5010,7 +5016,7 @@ async fn export_target_allowed( config: policy.config.clone(), }; let outcome = tokio::time::timeout( - Duration::from_millis(policy.timeout_millis), + export_activation_timeout(policy.timeout_millis), export_activation_policies.evaluate(&policy.provider, request), ) .await; @@ -5036,6 +5042,13 @@ async fn export_target_allowed( 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, + )) +} + const fn export_activation_target_kind_name(kind: ExportActivationTargetKind) -> &'static str { match kind { ExportActivationTargetKind::OtlpTrace => "otlp_trace", @@ -5052,7 +5065,7 @@ fn default_observability_config_version() -> u32 { } fn default_export_activation_timeout_millis() -> u64 { - 5_000 + DEFAULT_EXPORT_ACTIVATION_TIMEOUT_MILLIS } fn default_atof_mode() -> String { diff --git a/crates/core/tests/unit/observability/plugin_component_tests.rs b/crates/core/tests/unit/observability/plugin_component_tests.rs index 50b9c24d3..a89778766 100644 --- a/crates/core/tests/unit/observability/plugin_component_tests.rs +++ b/crates/core/tests/unit/observability/plugin_component_tests.rs @@ -5082,7 +5082,7 @@ async fn unavailable_and_timed_out_export_policies_fail_closed() { "test.slow-export-policy", Arc::new(|_| { Box::pin(async { - tokio::time::sleep(Duration::from_millis(25)).await; + tokio::time::sleep(Duration::from_millis(1_100)).await; Ok(ExportActivationDecision::Allow) }) }), @@ -5091,7 +5091,7 @@ async fn unavailable_and_timed_out_export_policies_fail_closed() { let mut timed_out = test_opentelemetry_endpoint(); timed_out.activation_policy = Some(ExportActivationPolicyConfig { provider: "test.slow-export-policy".into(), - timeout_millis: 1, + timeout_millis: 1_000, config: Json::Null, }); let mut context = PluginRegistrationContext::with_export_activation_policies( @@ -5264,14 +5264,24 @@ fn export_activation_policy_config_validates_provider_and_timeout() { "version": 4, "opentelemetry": { "enabled": true, - "endpoints": [{ - "type": "full", - "endpoint": "http://localhost:4318/v1/traces", - "activation_policy": { - "provider": " ", - "timeout_millis": 60001 + "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()); @@ -5282,4 +5292,23 @@ fn export_activation_policy_config_validates_provider_and_timeout() { 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/node/observability.js b/crates/node/observability.js index 6268fcfee..8e35b7bf5 100644 --- a/crates/node/observability.js +++ b/crates/node/observability.js @@ -60,12 +60,12 @@ function exportActivationPolicy(config) { 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 ?? 5000; - if (!Number.isInteger(timeout) || timeout < 1 || timeout > 60000) { - throw new TypeError('Export activation policy timeout_millis must be between 1 and 60000'); + 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 { - timeout_millis: 5000, + timeout_millis: 30000, config: null, ...config, }; diff --git a/crates/node/tests/observability_plugin_tests.mjs b/crates/node/tests/observability_plugin_tests.mjs index 42c3d604d..298aa7cdf 100644 --- a/crates/node/tests/observability_plugin_tests.mjs +++ b/crates/node/tests/observability_plugin_tests.mjs @@ -25,7 +25,7 @@ describe('observability plugin helpers', () => { }); assert.deepEqual(policy, { provider: 'com.example.runtime-policy', - timeout_millis: 5000, + timeout_millis: 30000, config: { allowed_countries: ['US', 'CA'] }, }); assert.equal( @@ -38,8 +38,12 @@ describe('observability plugin helpers', () => { ); assert.throws(() => observability.exportActivationPolicy({ provider: ' ' }), /nonblank/); assert.throws( - () => observability.exportActivationPolicy({ provider: 'test', timeout_millis: 60001 }), - /between 1 and 60000/, + () => observability.exportActivationPolicy({ provider: 'test', timeout_millis: 999 }), + /between 1000 and 300000/, + ); + assert.throws( + () => observability.exportActivationPolicy({ provider: 'test', timeout_millis: 300001 }), + /between 1000 and 300000/, ); }); diff --git a/docs/configure-plugins/observability/configuration.mdx b/docs/configure-plugins/observability/configuration.mdx index 430780039..b052ed70f 100644 --- a/docs/configure-plugins/observability/configuration.mdx +++ b/docs/configure-plugins/observability/configuration.mdx @@ -98,7 +98,7 @@ declares the `export_activation_policy` capability and registers one callback: ```toml [components.config.opentelemetry.endpoints.activation_policy] provider = "com.example.runtime-policy" -timeout_millis = 5000 +timeout_millis = 30000 config = { allowed_countries = ["US", "CA"], min_available_memory_mib = 2048 } ``` @@ -121,7 +121,10 @@ 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 5000 milliseconds and must be between 1 and 60000. +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, diff --git a/go/nemo_relay/observability_plugin.go b/go/nemo_relay/observability_plugin.go index ef3e79029..1142a8a85 100644 --- a/go/nemo_relay/observability_plugin.go +++ b/go/nemo_relay/observability_plugin.go @@ -263,7 +263,7 @@ func NewObservabilityConfig() ObservabilityConfig { // NewExportActivationPolicyConfig returns an activation policy with the canonical timeout. func NewExportActivationPolicyConfig(provider string) ExportActivationPolicyConfig { - return ExportActivationPolicyConfig{Provider: provider, TimeoutMillis: 5000} + return ExportActivationPolicyConfig{Provider: provider, TimeoutMillis: 30000} } // NewObservabilityAtofConfig returns disabled ATOF JSONL settings with native defaults. diff --git a/go/nemo_relay/observability_plugin_test.go b/go/nemo_relay/observability_plugin_test.go index f8ed333f6..ac48dda20 100644 --- a/go/nemo_relay/observability_plugin_test.go +++ b/go/nemo_relay/observability_plugin_test.go @@ -205,7 +205,7 @@ func assertWrappedObservabilityConfig(t *testing.T, wrapped PluginComponentSpec) 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(5000) { + if !ok || policy["provider"] != "com.example.runtime-policy" || policy["timeout_millis"] != float64(30000) { t.Fatalf("unexpected export activation policy: %#v", value) } } diff --git a/python/nemo_relay/observability.py b/python/nemo_relay/observability.py index e4c32fba4..04c501b49 100644 --- a/python/nemo_relay/observability.py +++ b/python/nemo_relay/observability.py @@ -57,7 +57,7 @@ class ExportActivationPolicyConfig: """Activation-time policy attached to one remote exporter target.""" provider: str - timeout_millis: int = 5000 + timeout_millis: int = 30000 config: Json = None def to_dict(self) -> JsonObject: diff --git a/python/tests/test_observability_plugin.py b/python/tests/test_observability_plugin.py index 270e387df..be3346cb2 100644 --- a/python/tests/test_observability_plugin.py +++ b/python/tests/test_observability_plugin.py @@ -95,7 +95,7 @@ def test_remote_targets_serialize_export_activation_policy(self): ) expected = { "provider": "com.example.runtime-policy", - "timeout_millis": 5000, + "timeout_millis": 30000, "config": {"allowed_countries": ["US", "CA"]}, } assert ( From 1878b8f74dbf6b8ebea91100e062a9f778d3bdae Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 20 Aug 2026 09:51:46 -0700 Subject: [PATCH 5/9] fix(exporters): preserve configured target indices Signed-off-by: Bryan Bednarski --- .../src/observability/plugin_component.rs | 67 +++++++++++++------ .../observability/plugin_component_tests.rs | 57 ++++++++++++++++ 2 files changed, 104 insertions(+), 20 deletions(-) diff --git a/crates/core/src/observability/plugin_component.rs b/crates/core/src/observability/plugin_component.rs index 862bf6433..1298fa94c 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -1162,7 +1162,7 @@ fn register_full_payload_policy( } async fn register_atof_exporter( - mut section: AtofSectionConfig, + section: AtofSectionConfig, ctx: &mut PluginRegistrationContext, ) -> PluginResult<()> { let export_activation_policies = ctx.export_activation_policies(); @@ -1181,17 +1181,14 @@ async fn register_atof_exporter( } }; if allowed { - allowed_sinks.push(sink); + allowed_sinks.push((index, sink)); } } - section.sinks = allowed_sinks; - if section.sinks.is_empty() { + if allowed_sinks.is_empty() { return Ok(()); } - let exporters = section - .sinks + let exporters = allowed_sinks .into_iter() - .enumerate() .map(|(index, sink)| { let config = CoreAtofExporterConfig { sink: build_atof_sink_config(index, sink)?, @@ -1306,21 +1303,32 @@ async fn register_atif_dispatcher( ) .await { - allowed_storage.push(entry); + allowed_storage.push((index, entry)); } } - section.storage = allowed_storage; - if had_remote_storage && section.storage.is_empty() { + if had_remote_storage && allowed_storage.is_empty() { return Ok(()); } - let mut storage_vec = Vec::with_capacity(section.storage.len()); - for (index, entry) in section.storage.iter().enumerate() { + 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, + ))); let dispatcher = atif_dispatcher_subscriber( Arc::clone(&manager), ctx.qualify_name("atif-"), @@ -2261,6 +2269,7 @@ fn shutdown_opentelemetry_providers( struct AtifDispatcher { config: AtifSectionConfig, + remote_storage_indices: Vec, agents: HashMap, scope_owners: HashMap, scope_subscribers: HashMap, @@ -2362,8 +2371,17 @@ enum SinkLabel { impl AtifDispatcher { fn new(config: AtifSectionConfig) -> Self { + let remote_storage_indices = (0..config.storage.len()).collect(); + Self::with_remote_storage_indices(config, remote_storage_indices) + } + + fn with_remote_storage_indices( + config: AtifSectionConfig, + remote_storage_indices: Vec, + ) -> Self { Self { config, + remote_storage_indices, agents: HashMap::new(), scope_owners: HashMap::new(), scope_subscribers: HashMap::new(), @@ -2658,7 +2676,9 @@ impl AtifDispatcher { if self.config.storage.is_empty() { vec![SinkLabel::Local] } else { - (0..self.config.storage.len()) + self.remote_storage_indices + .iter() + .copied() .map(SinkLabel::Remote) .collect() } @@ -2979,6 +2999,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| { @@ -2989,7 +3010,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) }) @@ -3021,19 +3046,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( diff --git a/crates/core/tests/unit/observability/plugin_component_tests.rs b/crates/core/tests/unit/observability/plugin_component_tests.rs index a89778766..943d43d30 100644 --- a/crates/core/tests/unit/observability/plugin_component_tests.rs +++ b/crates/core/tests/unit/observability/plugin_component_tests.rs @@ -5258,6 +5258,63 @@ async fn denied_atof_stream_retains_allowed_local_sink() { rollback_registrations(&mut registrations); } +#[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], + ); + assert_eq!(dispatcher.sink_targets(), vec![SinkLabel::Remote(2)]); +} + #[test] fn export_activation_policy_config_validates_provider_and_timeout() { let config = json!({ From fe59ee4c4fd97ac16fa961dd866bf3d63d11c849 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 20 Aug 2026 09:52:12 -0700 Subject: [PATCH 6/9] fix(node): retain export policy defaults Signed-off-by: Bryan Bednarski --- crates/node/observability.js | 4 ++-- crates/node/tests/observability_plugin_tests.mjs | 12 ++++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/crates/node/observability.js b/crates/node/observability.js index 8e35b7bf5..7ae92fafb 100644 --- a/crates/node/observability.js +++ b/crates/node/observability.js @@ -65,9 +65,9 @@ function exportActivationPolicy(config) { throw new TypeError('Export activation policy timeout_millis must be between 1000 and 300000'); } return { - timeout_millis: 30000, - config: null, ...config, + timeout_millis: timeout, + config: config.config ?? null, }; } diff --git a/crates/node/tests/observability_plugin_tests.mjs b/crates/node/tests/observability_plugin_tests.mjs index 298aa7cdf..06005f509 100644 --- a/crates/node/tests/observability_plugin_tests.mjs +++ b/crates/node/tests/observability_plugin_tests.mjs @@ -28,6 +28,18 @@ describe('observability plugin helpers', () => { 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', From 0a1f79012de8d1edbcecb5d0afc0558e70eed1ff Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 20 Aug 2026 09:53:02 -0700 Subject: [PATCH 7/9] test(exporters): verify mixed ATOF policy filtering Signed-off-by: Bryan Bednarski --- .../observability/plugin_component_tests.rs | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/crates/core/tests/unit/observability/plugin_component_tests.rs b/crates/core/tests/unit/observability/plugin_component_tests.rs index 943d43d30..30eb78a57 100644 --- a/crates/core/tests/unit/observability/plugin_component_tests.rs +++ b/crates/core/tests/unit/observability/plugin_component_tests.rs @@ -5227,18 +5227,21 @@ async fn all_denied_atif_storage_does_not_fall_back_to_local_files() { #[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(temp_dir("allowed-local-atof")), + output_directory: Some(output_directory.clone()), filename: Some("events.jsonl".into()), mode: "overwrite".into(), }), AtofSinkSectionConfig::Stream(AtofStreamSinkSectionConfig { name: Some("denied".into()), url: "https://collector.example/events".into(), - transport: "http_post".into(), + // Construction would fail if policy filtering retained this sink. + transport: "invalid".into(), headers: HashMap::new(), header_env: HashMap::new(), timeout_millis: 3_000, @@ -5254,8 +5257,20 @@ async fn denied_atof_stream_retains_allowed_local_sink() { let mut context = PluginRegistrationContext::new(); register_atof_exporter(section, &mut context).await.unwrap(); let mut registrations = context.into_registrations(); - assert!(!registrations.is_empty()); + 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] From da45b9ee5ee01bf88529ef4e33e2bc6cc08d141b Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 20 Aug 2026 09:53:55 -0700 Subject: [PATCH 8/9] test(plugins): cover native export policy outcomes Signed-off-by: Bryan Bednarski --- crates/plugin/tests/typed_callbacks.rs | 40 ++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index b472ae8bd..a3fb2a6f0 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -3568,6 +3568,46 @@ 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_async_registration(NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy); + let result = invoke_async_registration( + &host, + &deny, + json!({"target_kind": "atif_s3", "config": {"enabled": false}}), + None, + ) + .unwrap(); + assert_eq!(result, json!("deny")); + unsafe { deny.free() }; + + let failure = + take_async_registration(NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy); + let error = invoke_async_registration( + &host, + &failure, + json!({"target_kind": "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(); From a59fc8b9b68a7faefbbc64e7437a719126682249 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 20 Aug 2026 12:54:49 -0700 Subject: [PATCH 9/9] feat(exporters): add public export activation hooks Signed-off-by: Bryan Bednarski --- crates/core/src/api/export_activation.rs | 285 +++++++++++++++++- .../src/observability/plugin_component.rs | 102 ++++--- crates/core/src/plugin.rs | 89 +++++- crates/core/src/plugin/dynamic.rs | 2 +- crates/core/src/plugin/dynamic/native.rs | 80 ++++- crates/core/src/plugin/dynamic/worker.rs | 103 ++++++- .../core/tests/unit/dynamic_worker_tests.rs | 42 ++- crates/core/tests/unit/native_plugin_tests.rs | 5 + .../observability/plugin_component_tests.rs | 66 +++- crates/core/tests/unit/plugin_tests.rs | 1 + crates/ffi/nemo_relay.h | 34 +++ crates/ffi/src/api/mod.rs | 23 +- crates/ffi/src/api/plugin.rs | 88 +++++- crates/ffi/src/callable.rs | 67 ++++ crates/node/observability.d.ts | 4 +- crates/node/observability.js | 2 +- crates/node/plugin.d.ts | 29 ++ crates/node/src/api/mod.rs | 129 +++++++- .../node/tests/observability_plugin_tests.mjs | 6 + .../tests/plugin_export_activation_tests.mjs | 56 ++++ crates/plugin/src/async_sdk.rs | 62 +++- crates/plugin/src/lib.rs | 67 +++- crates/plugin/tests/typed_callbacks.rs | 73 +++-- crates/plugin/tests/unit/async_sdk_tests.rs | 4 - crates/python/src/py_callable.rs | 71 +++++ crates/python/src/py_plugin.rs | 96 +++++- .../coverage/py_plugin_coverage_tests.rs | 14 + crates/types/src/plugin.rs | 135 ++++++++- .../nemo/relay/worker/v1/plugin_worker.proto | 3 + crates/worker/README.md | 2 +- crates/worker/src/lib.rs | 62 +++- crates/worker/tests/worker_sdk_tests.rs | 44 ++- .../grpc-worker/python/about.mdx | 37 ++- .../rust-native-plugin-example.mdx | 40 ++- .../language-binding/register-behavior.mdx | 51 ++++ docs/configure-plugins/observability/atif.mdx | 7 + docs/configure-plugins/observability/atof.mdx | 5 + .../observability/configuration.mdx | 7 +- .../observability/opentelemetry.mdx | 6 + examples/python-grpc-worker-plugin/README.md | 2 +- examples/rust-native-plugin/README.md | 2 +- go/nemo_relay/callbacks.go | 33 ++ go/nemo_relay/observability_plugin.go | 28 +- go/nemo_relay/observability_plugin_test.go | 11 +- go/nemo_relay/plugin.go | 69 +++++ go/nemo_relay/plugin_gap_test.go | 67 ++++ python/nemo_relay/_native.pyi | 2 + python/nemo_relay/observability.py | 6 +- python/nemo_relay/observability.pyi | 2 + python/nemo_relay/plugin.py | 46 ++- python/nemo_relay/plugin.pyi | 24 +- python/plugin/README.md | 2 +- .../plugin/src/nemo_relay_plugin/__init__.py | 15 +- python/plugin/src/nemo_relay_plugin/_api.py | 102 ++++++- .../plugin/test_public_api_docstrings.py | 1 + python/tests/plugin/test_worker_sdk.py | 28 +- python/tests/test_observability_plugin.py | 2 + python/tests/test_plugin_export_activation.py | 57 ++++ 58 files changed, 2249 insertions(+), 249 deletions(-) create mode 100644 crates/node/tests/plugin_export_activation_tests.mjs create mode 100644 python/tests/test_plugin_export_activation.py diff --git a/crates/core/src/api/export_activation.rs b/crates/core/src/api/export_activation.rs index a5bc9b35b..e44834dbd 100644 --- a/crates/core/src/api/export_activation.rs +++ b/crates/core/src/api/export_activation.rs @@ -1,16 +1,19 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Activation-time policy hooks for Relay-managed remote exporters. +//! Activation-time hooks for plugin-managed export targets. use std::collections::HashMap; use std::future::Future; use std::pin::Pin; -use std::sync::{Arc, RwLock}; +use std::sync::{Arc, Mutex, RwLock}; +use std::time::Duration; use crate::error::{FlowError, Result}; pub use nemo_relay_types::plugin::{ - ExportActivationDecision, ExportActivationRequest, ExportActivationTargetKind, + ExportActivationDecision, ExportActivationPolicyConfig, ExportActivationRequest, + ExportActivationTargetKind, ExportTargetRegistration, MAX_EXPORT_ACTIVATION_TIMEOUT_MILLIS, + MIN_EXPORT_ACTIVATION_TIMEOUT_MILLIS, }; /// Asynchronous callback registered by one export-activation policy provider. @@ -22,18 +25,31 @@ pub type ExportActivationPolicyFn = Arc< + Sync, >; -/// Export-activation policy callbacks owned by one plugin activation. +/// 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)] -pub(crate) struct ExportActivationPolicyRegistry { +#[doc(hidden)] +pub struct ExportActivationRegistry { policies: RwLock>, + targets: Mutex>, + active_targets: Mutex>, } -impl ExportActivationPolicyRegistry { - pub(crate) fn register( - &self, - provider: &str, - callback: ExportActivationPolicyFn, - ) -> Result<()> { +#[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}" @@ -46,7 +62,8 @@ impl ExportActivationPolicyRegistry { Ok(()) } - pub(crate) fn deregister(&self, provider: &str) -> Result { + /// 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()) @@ -75,4 +92,248 @@ impl ExportActivationPolicyRegistry { .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/observability/plugin_component.rs b/crates/core/src/observability/plugin_component.rs index 1298fa94c..64b3e73f6 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -37,8 +37,9 @@ use uuid::Uuid; use crate::api::event::{Event, LogSeverity, ScopeCategory, ValidatedMetricMeasurement}; use crate::api::export_activation::{ - ExportActivationDecision, ExportActivationPolicyRegistry, ExportActivationRequest, - ExportActivationTargetKind, + 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; @@ -193,24 +194,6 @@ pub struct OpenTelemetrySectionConfig { pub metrics: Option, } -/// Activation-time policy attached to one remote exporter target. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -pub struct ExportActivationPolicyConfig { - /// Dynamic plugin identifier that owns the policy 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, -} - -const MIN_EXPORT_ACTIVATION_TIMEOUT_MILLIS: u64 = 1_000; -const DEFAULT_EXPORT_ACTIVATION_TIMEOUT_MILLIS: u64 = 30_000; -const MAX_EXPORT_ACTIVATION_TIMEOUT_MILLIS: u64 = 300_000; - /// Signal-common OTLP destination fields used by logs and metrics. #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] @@ -434,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. @@ -509,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 @@ -532,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(), } } @@ -897,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 }, } } @@ -923,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 }, } } @@ -1169,7 +1162,15 @@ async fn register_atof_exporter( 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(_) => true, + 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, @@ -1283,6 +1284,13 @@ async fn register_atif_dispatcher( .map_err(PluginError::InvalidConfig)?; 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 { @@ -1306,7 +1314,8 @@ async fn register_atif_dispatcher( allowed_storage.push((index, entry)); } } - if had_remote_storage && allowed_storage.is_empty() { + if (had_remote_storage && allowed_storage.is_empty()) || (!had_remote_storage && !local_allowed) + { return Ok(()); } @@ -1328,6 +1337,7 @@ async fn register_atif_dispatcher( 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), @@ -2270,6 +2280,7 @@ fn shutdown_opentelemetry_providers( struct AtifDispatcher { config: AtifSectionConfig, remote_storage_indices: Vec, + local_allowed: bool, agents: HashMap, scope_owners: HashMap, scope_subscribers: HashMap, @@ -2370,18 +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) + 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(), @@ -2669,12 +2683,15 @@ 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 { self.remote_storage_indices .iter() @@ -3024,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), @@ -3346,6 +3364,7 @@ fn validate_observability_section_fields( "extra", "output_directory", "filename_template", + "local_activation_policy", "storage", ], ); @@ -4423,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, @@ -4700,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, @@ -5039,7 +5070,7 @@ async fn export_target_allowed( return true; }; let request = ExportActivationRequest { - target_kind, + target_kind: target_kind.clone(), config: policy.config.clone(), }; let outcome = tokio::time::timeout( @@ -5060,10 +5091,10 @@ async fn export_target_allowed( event = "export_activation_policy_denied", plugin_kind = OBSERVABILITY_PLUGIN_KIND, provider = policy.provider.as_str(), - target_kind = export_activation_target_kind_name(target_kind), + target_kind = target_kind.as_str(), field, reason; - "Remote exporter target suppressed by activation policy" + "Exporter target suppressed by activation policy" ); } allowed @@ -5076,25 +5107,10 @@ fn export_activation_timeout(timeout_millis: u64) -> Duration { )) } -const fn export_activation_target_kind_name(kind: ExportActivationTargetKind) -> &'static str { - match kind { - ExportActivationTargetKind::OtlpTrace => "otlp_trace", - ExportActivationTargetKind::OtlpLog => "otlp_log", - ExportActivationTargetKind::OtlpMetric => "otlp_metric", - ExportActivationTargetKind::AtofStream => "atof_stream", - ExportActivationTargetKind::AtifHttp => "atif_http", - ExportActivationTargetKind::AtifS3 => "atif_s3", - } -} - fn default_observability_config_version() -> u32 { 4 } -fn default_export_activation_timeout_millis() -> u64 { - DEFAULT_EXPORT_ACTIVATION_TIMEOUT_MILLIS -} - fn default_atof_mode() -> String { "append".to_string() } diff --git a/crates/core/src/plugin.rs b/crates/core/src/plugin.rs index b0cba4032..3006cf49f 100644 --- a/crates/core/src/plugin.rs +++ b/crates/core/src/plugin.rs @@ -22,7 +22,10 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value as Json}; use thiserror::Error; -use crate::api::export_activation::{ExportActivationPolicyFn, ExportActivationPolicyRegistry}; +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, @@ -405,6 +408,7 @@ impl PluginRegistration { pub struct PluginRegistrationContext { registrations: Vec, namespace: Option, + provider_id: Option, export_activation_policies: Arc, } @@ -413,6 +417,7 @@ impl Default for PluginRegistrationContext { Self { registrations: Vec::new(), namespace: None, + provider_id: None, export_activation_policies: Arc::new(ExportActivationPolicyRegistry::default()), } } @@ -429,17 +434,20 @@ 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, } } @@ -451,14 +459,21 @@ impl PluginRegistrationContext { Self { registrations: Vec::new(), namespace: None, + provider_id: None, export_activation_policies, } } - pub(crate) fn export_activation_policies(&self) -> Arc { + #[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 @@ -472,7 +487,7 @@ impl PluginRegistrationContext { } /// Registers the single export-activation policy owned by `provider`. - pub fn register_export_activation_policy( + pub(crate) fn register_export_activation_policy_for_provider( &mut self, provider: &str, callback: ExportActivationPolicyFn, @@ -502,6 +517,53 @@ impl PluginRegistrationContext { 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); @@ -2729,10 +2791,17 @@ async fn initialize_plugin_components( 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!( @@ -2753,6 +2822,7 @@ async fn initialize_plugin_components( let mut pending = PendingPluginRegistrationContext::new( namespace, + component.kind.clone(), rollback_failures.clone(), Arc::clone(&export_activation_policies), ); @@ -2762,6 +2832,13 @@ async fn initialize_plugin_components( registrations.extend(pending.take()); } + export_activation_policies + .activate_targets() + .await + .map_err(|error| { + PluginError::RegistrationFailed(format!("export target activation failed: {error}")) + })?; + Ok(registrations.take()) } @@ -2804,12 +2881,14 @@ struct PendingPluginRegistrationContext { impl PendingPluginRegistrationContext { fn new( namespace: String, + provider_id: String, rollback_failures: Option>>>, export_activation_policies: Arc, ) -> Self { Self { context: PluginRegistrationContext::with_namespace_and_export_activation_policies( namespace, + provider_id, export_activation_policies, ), rollback_failures, diff --git a/crates/core/src/plugin/dynamic.rs b/crates/core/src/plugin/dynamic.rs index 61e5fc799..9add23fa7 100644 --- a/crates/core/src/plugin/dynamic.rs +++ b/crates/core/src/plugin/dynamic.rs @@ -244,7 +244,7 @@ pub enum DynamicPluginCapability { PluginWorker, /// Typed configuration schema contribution capability. ConfigSchema, - /// Activation-time policy for Relay-managed remote exporters. + /// Activation-time policy for local or remote export targets. ExportActivationPolicy, } diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 886cecb52..bae5a268b 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -24,6 +24,7 @@ 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::{ @@ -56,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, @@ -80,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, }; @@ -285,6 +287,7 @@ struct NativePluginInstance { plugin_kind: String, relay_compat: String, allows_multiple_components: bool, + allows_export_activation_policy: bool, plugin: Mutex, _library: Library, } @@ -445,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, })) @@ -3540,6 +3547,23 @@ fn wrap_native_async_export_activation_policy( }) } +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, @@ -3736,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(()) => { @@ -3849,13 +3912,6 @@ unsafe extern "C" fn native_plugin_context_register_async_middleware( priority, wrap_native_async_event_metadata_injector(instance, cb, user_data, free_fn), ), - NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy => { - let provider = instance.plugin_kind.clone(); - context.register_export_activation_policy( - &provider, - wrap_native_async_export_activation_policy(instance, cb, user_data, free_fn), - ) - } }; match registration { Ok(()) => NemoRelayStatus::Ok, diff --git a/crates/core/src/plugin/dynamic/worker.rs b/crates/core/src/plugin/dynamic/worker.rs index 661d5386e..421a464ce 100644 --- a/crates/core/src/plugin/dynamic/worker.rs +++ b/crates/core/src/plugin/dynamic/worker.rs @@ -66,6 +66,7 @@ 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::{ @@ -91,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, }; @@ -100,6 +101,7 @@ 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); @@ -620,8 +622,12 @@ fn load_one_worker_plugin( return Err(worker_error_to_plugin(error, "worker registration failed")); } validate_registration_plan(&spec.plugin_id, ®ister)?; - filter_unadvertised_export_activation_policy( + filter_unadvertised_export_activation_hooks( &spec.plugin_id, + manifest + .capabilities + .items + .contains(&DynamicPluginCapability::ExportActivationPolicy), &handshake.supported_surfaces, &mut register.registrations, ); @@ -1082,6 +1088,9 @@ impl WorkerPluginInstance { )?, 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 @@ -1161,7 +1170,6 @@ impl WorkerPluginInstance { ) -> crate::plugin::Result<()> { let callback = Arc::new(self.clone_for_callback()); let callback_name = name.to_owned(); - let provider = self.plugin_kind.clone(); let policy: ExportActivationPolicyFn = Arc::new(move |request| { let callback = Arc::clone(&callback); let callback_name = callback_name.clone(); @@ -1171,7 +1179,40 @@ impl WorkerPluginInstance { .await }) }); - ctx.register_export_activation_policy(&provider, policy) + 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( @@ -1601,6 +1642,25 @@ impl WorkerPluginCallback { }) } + 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, @@ -3396,25 +3456,44 @@ fn validate_registration_plan( Ok(()) } -fn filter_unadvertised_export_activation_policy( +fn filter_unadvertised_export_activation_hooks( plugin_id: &str, + policy_capability_declared: bool, supported_surfaces: &[i32], registrations: &mut Vec, ) { - let policy_surface = RegistrationSurface::ExportActivationPolicy as i32; - if supported_surfaces.contains(&policy_surface) { - return; - } registrations.retain(|registration| { - if registration.surface != policy_surface { + 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 = "export_activation_policy"; - "Worker export activation policy registration ignored because the surface was not advertised" + surface = surface_name; + "Worker export activation hook registration ignored because the surface was not advertised" ); false }); diff --git a/crates/core/tests/unit/dynamic_worker_tests.rs b/crates/core/tests/unit/dynamic_worker_tests.rs index 06815fe13..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, }, @@ -562,7 +565,7 @@ fn registration_plan_and_scope_type_helpers_validate_edges() { } #[test] -fn unadvertised_export_activation_policy_registration_is_ignored() { +fn unadvertised_export_activation_hook_registration_is_ignored() { let policy = registration( RegistrationSurface::ExportActivationPolicy, "export_activation_policy", @@ -570,20 +573,50 @@ fn unadvertised_export_activation_policy_registration_is_ignored() { let subscriber = registration(RegistrationSurface::Subscriber, "subscriber"); let mut registrations = vec![policy.clone(), subscriber.clone()]; - filter_unadvertised_export_activation_policy( + 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_policy( + filter_unadvertised_export_activation_hooks( "fixture_worker", + true, &[RegistrationSurface::ExportActivationPolicy as i32], &mut advertised, ); - assert_eq!(advertised, vec![policy]); + 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] @@ -2917,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 30eb78a57..80f559060 100644 --- a/crates/core/tests/unit/observability/plugin_component_tests.rs +++ b/crates/core/tests/unit/observability/plugin_component_tests.rs @@ -5010,7 +5010,7 @@ async fn export_activation_policy_is_evaluated_once_before_otel_construction() { let callback_calls = Arc::clone(&calls); let mut provider_context = PluginRegistrationContext::new(); provider_context - .register_export_activation_policy( + .register_export_activation_policy_for_provider( "test.allow-export", Arc::new(move |request| { let callback_calls = Arc::clone(&callback_calls); @@ -5078,7 +5078,7 @@ async fn unavailable_and_timed_out_export_policies_fail_closed() { let mut provider_context = PluginRegistrationContext::new(); provider_context - .register_export_activation_policy( + .register_export_activation_policy_for_provider( "test.slow-export-policy", Arc::new(|_| { Box::pin(async { @@ -5130,7 +5130,7 @@ async fn denied_and_failed_export_policies_fail_closed() { let mut provider_context = PluginRegistrationContext::new(); let callback_result = result.clone(); provider_context - .register_export_activation_policy( + .register_export_activation_policy_for_provider( provider, Arc::new(move |_| { let callback_result = callback_result.clone(); @@ -5171,24 +5171,27 @@ fn export_activation_provider_uniqueness_is_scoped_to_one_activation() { let callback = Arc::new(|_| Box::pin(async { Ok(ExportActivationDecision::Allow) }) as _); let mut first = PluginRegistrationContext::new(); first - .register_export_activation_policy("test.duplicate-policy", callback.clone()) + .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("test.duplicate-policy", callback.clone()) + .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("test.duplicate-policy", callback.clone()) + .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("test.duplicate-policy", callback) + .register_export_activation_policy_for_provider("test.duplicate-policy", callback) .unwrap(); let mut duplicate_registrations = duplicate.into_registrations(); rollback_registrations(&mut duplicate_registrations); @@ -5236,6 +5239,7 @@ async fn denied_atof_stream_retains_allowed_local_sink() { output_directory: Some(output_directory.clone()), filename: Some("events.jsonl".into()), mode: "overwrite".into(), + activation_policy: None, }), AtofSinkSectionConfig::Stream(AtofStreamSinkSectionConfig { name: Some("denied".into()), @@ -5273,6 +5277,53 @@ async fn denied_atof_stream_retains_allowed_local_sink() { 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() { @@ -5326,6 +5377,7 @@ fn filtered_atif_storage_preserves_original_sink_labels() { ..AtifSectionConfig::default() }, vec![2], + true, ); assert_eq!(dispatcher.sink_targets(), vec![SinkLabel::Remote(2)]); } diff --git a/crates/core/tests/unit/plugin_tests.rs b/crates/core/tests/unit/plugin_tests.rs index 59b12c916..cffe24a88 100644 --- a/crates/core/tests/unit/plugin_tests.rs +++ b/crates/core/tests/unit/plugin_tests.rs @@ -1737,6 +1737,7 @@ fn test_pending_rollbacks_ignore_delivery_only_errors() { { let mut pending = PendingPluginRegistrationContext::new( "fixture.".into(), + "fixture".into(), Some(Arc::clone(&failures)), Arc::new(ExportActivationPolicyRegistry::default()), ); 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 81ec4444a..eebe0b260 100644 --- a/crates/node/observability.d.ts +++ b/crates/node/observability.d.ts @@ -24,6 +24,7 @@ export interface AtofFileSinkConfig { output_directory?: string; filename?: string; mode?: 'append' | 'overwrite' | string; + activation_policy?: ExportActivationPolicyConfig; } export interface AtofStreamSinkConfig { @@ -73,6 +74,7 @@ export interface AtifConfig { output_directory?: string; filename_template?: string; storage?: S3StorageConfig | HttpStorageConfig | Array; + local_activation_policy?: ExportActivationPolicyConfig; } export interface OpenTelemetryEndpointConfig { @@ -160,7 +162,7 @@ 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 remote-export activation policy attachment. */ +/** Create and validate one export-activation policy attachment. */ export declare function exportActivationPolicy( config: ExportActivationPolicyConfig, ): ExportActivationPolicyConfig; diff --git a/crates/node/observability.js b/crates/node/observability.js index 7ae92fafb..4fe93d5d8 100644 --- a/crates/node/observability.js +++ b/crates/node/observability.js @@ -48,7 +48,7 @@ function atifConfig(config = {}) { } /** - * Create and validate one remote-export activation policy attachment. + * 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. 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 06005f509..54aa689f3 100644 --- a/crates/node/tests/observability_plugin_tests.mjs +++ b/crates/node/tests/observability_plugin_tests.mjs @@ -48,6 +48,12 @@ describe('observability plugin helpers', () => { }).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 }), 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/plugin/src/async_sdk.rs b/crates/plugin/src/async_sdk.rs index 27febd116..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,18 +1192,16 @@ impl PluginContext<'_> { ) } - /// Registers this plugin's activation-time policy for remote exporters. + /// 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_unary_adapter( - NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy, + self.register_activation_adapter( + NemoRelayNativeActivationHookKind::ExportActivationPolicy, "export_activation_policy", - 0, - false, Box::new(move |value, _, _| { let callback = Arc::clone(&callback); Box::pin(async move { @@ -1191,6 +1216,32 @@ impl PluginContext<'_> { ) } + /// 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, @@ -1635,7 +1686,6 @@ fn registration_operation(kind: NemoRelayNativeAsyncMiddlewareKind) -> &'static NemoRelayNativeAsyncMiddlewareKind::ScopeSanitizeStart => "scope start sanitizer", NemoRelayNativeAsyncMiddlewareKind::ScopeSanitizeEnd => "scope end sanitizer", NemoRelayNativeAsyncMiddlewareKind::EventMetadataInjector => "Event metadata injector", - NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy => "export activation policy", } } diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 9334f4802..8873889a4 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -41,8 +41,8 @@ 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, ExportActivationDecision, ExportActivationRequest, - ExportActivationTargetKind, + ConfigDiagnostic, DiagnosticLevel, ExportActivationDecision, ExportActivationPolicyConfig, + ExportActivationRequest, ExportActivationTargetKind, ExportTargetRegistration, }; use serde::{Serialize, de::DeserializeOwned}; use serde_json::Map; @@ -831,8 +831,6 @@ pub enum NemoRelayNativeAsyncMiddlewareKind { ScopeSanitizeEnd = 13, /// Event metadata injector. EventMetadataInjector = 14, - /// Activation-time policy for Relay-managed remote exporters. - ExportActivationPolicy = 15, } impl TryFrom for NemoRelayNativeAsyncMiddlewareKind { @@ -855,7 +853,28 @@ impl TryFrom for NemoRelayNativeAsyncMiddlewareKind { 12 => Ok(Self::ScopeSanitizeStart), 13 => Ok(Self::ScopeSanitizeEnd), 14 => Ok(Self::EventMetadataInjector), + _ => Err(()), + } + } +} + +/// 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(()), } } @@ -2480,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 a3fb2a6f0..846c01f50 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -23,11 +23,12 @@ use futures::StreamExt; use nemo_relay_plugin::{ AnnotatedLlmRequest, BuiltinLlmCodec, CategoryProfile, ConfigDiagnostic, DataSchema, DiagnosticLevel, Event, EventCategory, EventSanitizeFields, ExportActivationDecision, - ExportActivationTargetKind, 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, + 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, + NemoRelayNativeActivationHookKind, NemoRelayNativeAsyncCallbackState, NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncLlmStreamOpenCb, NemoRelayNativeAsyncLlmStreamPullCb, NemoRelayNativeAsyncMiddlewareCb, NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncNextResultCb, @@ -69,13 +70,17 @@ fn async_abi_discriminants_reject_unknown_values() { Kind::ScopeSanitizeStart, Kind::ScopeSanitizeEnd, Kind::EventMetadataInjector, - Kind::ExportActivationPolicy, ]; for (discriminant, kind) in middleware_kinds.into_iter().enumerate() { assert_eq!(kind as u32, discriminant as u32); assert_eq!(Kind::try_from(discriminant as u32), Ok(kind)); } - assert!(NemoRelayNativeAsyncMiddlewareKind::try_from(16).is_err()); + 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) @@ -288,7 +293,7 @@ struct RegisteredLlmRequestIntercept { } struct RegisteredAsync { - kind: NemoRelayNativeAsyncMiddlewareKind, + kind: u32, name: String, priority: i32, break_chain: bool, @@ -1856,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, @@ -2169,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) } @@ -2184,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(); @@ -3303,6 +3312,15 @@ fn typed_async_middleware_registers_and_round_trips_every_surface() { 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() @@ -3317,17 +3335,17 @@ fn typed_async_middleware_registers_and_round_trips_every_surface() { ) }) .collect::>(); - assert_eq!(metadata.len(), 15); + 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); { @@ -3362,17 +3380,25 @@ fn typed_async_middleware_registers_and_round_trips_every_surface() { unsafe { metadata_registration.free() }; let policy_registration = - take_async_registration(NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy); + take_activation_registration(NemoRelayNativeActivationHookKind::ExportActivationPolicy); let result = invoke_async_registration( &host, &policy_registration, - json!({"target_kind": "otlp_trace", "config": {"enabled": true}}), + 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, @@ -3583,11 +3609,12 @@ fn typed_export_activation_policy_round_trips_deny_and_error() { ctx.register_export_activation_policy(|_| async move { Err("policy failed".into()) }) .unwrap(); - let deny = take_async_registration(NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy); + let deny = + take_activation_registration(NemoRelayNativeActivationHookKind::ExportActivationPolicy); let result = invoke_async_registration( &host, &deny, - json!({"target_kind": "atif_s3", "config": {"enabled": false}}), + json!({"target_kind": "nemo_relay.atif.s3", "config": {"enabled": false}}), None, ) .unwrap(); @@ -3595,11 +3622,11 @@ fn typed_export_activation_policy_round_trips_deny_and_error() { unsafe { deny.free() }; let failure = - take_async_registration(NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy); + take_activation_registration(NemoRelayNativeActivationHookKind::ExportActivationPolicy); let error = invoke_async_registration( &host, &failure, - json!({"target_kind": "otlp_trace", "config": null}), + json!({"target_kind": "nemo_relay.otlp.trace", "config": null}), None, ) .unwrap_err(); diff --git a/crates/plugin/tests/unit/async_sdk_tests.rs b/crates/plugin/tests/unit/async_sdk_tests.rs index 234c67319..4d389674e 100644 --- a/crates/plugin/tests/unit/async_sdk_tests.rs +++ b/crates/plugin/tests/unit/async_sdk_tests.rs @@ -125,10 +125,6 @@ fn async_sdk_status_helpers_keep_operation_labels_stable() { NemoRelayNativeAsyncMiddlewareKind::EventMetadataInjector, "Event metadata injector", ), - ( - NemoRelayNativeAsyncMiddlewareKind::ExportActivationPolicy, - "export activation policy", - ), ]; for (kind, expected) in labels { assert_eq!(registration_operation(kind), expected); 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 6c66e73ec..22ed4d7c7 100644 --- a/crates/types/src/plugin.rs +++ b/crates/types/src/plugin.rs @@ -3,34 +3,103 @@ //! Shared plugin diagnostic data types. -use serde::{Deserialize, Serialize}; +use std::borrow::Cow; + +use serde::{Deserialize, Deserializer, Serialize}; use crate::Json; -/// Remote exporter target presented to an activation policy. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +/// Export target kind presented to an activation policy. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum ExportActivationTargetKind { +#[serde(transparent)] +pub struct ExportActivationTargetKind(Cow<'static, str>); + +impl ExportActivationTargetKind { /// OpenTelemetry trace destination. - OtlpTrace, + pub const OTLP_TRACE: Self = Self(Cow::Borrowed("nemo_relay.otlp.trace")); /// OpenTelemetry log destination. - OtlpLog, + pub const OTLP_LOG: Self = Self(Cow::Borrowed("nemo_relay.otlp.log")); /// OpenTelemetry metric destination. - OtlpMetric, + 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. - AtofStream, + 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. - AtifHttp, + pub const ATIF_HTTP: Self = Self(Cow::Borrowed("nemo_relay.atif.http")); /// Remote ATIF S3-compatible storage. - AtifS3, + 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 remote exporter target being considered. + /// Kind of local or remote export target being considered. pub target_kind: ExportActivationTargetKind, /// Opaque target-local policy configuration. #[serde(default)] @@ -42,12 +111,50 @@ pub struct ExportActivationRequest { #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] #[serde(rename_all = "snake_case")] pub enum ExportActivationDecision { - /// Construct and activate the remote exporter target. + /// Construct and activate the exporter target. Allow, - /// Suppress the remote exporter target for this activation. + /// 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)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] 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 23146dbf3..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 @@ -46,6 +46,7 @@ enum RegistrationSurface { 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; @@ -151,6 +152,7 @@ message Registration { int32 priority = 3; bool break_chain = 4; reserved 5; + JsonEnvelope export_target = 6; } message InvokeRequest { @@ -167,6 +169,7 @@ message InvokeRequest { 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 496662d41..b710fd534 100644 --- a/crates/worker/README.md +++ b/crates/worker/README.md @@ -42,7 +42,7 @@ outcome fields now use structural protobuf messages. - **`WorkerPlugin`**: The plugin identity, validation, and registration contract. - **`PluginContext`**: Typed registrations for all supported worker surfaces, - including activation-time policies for Relay-managed remote exporters. + 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 0e7ae87ac..cf4bb47f7 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -52,8 +52,8 @@ 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, ExportActivationDecision, ExportActivationRequest, - ExportActivationTargetKind, + 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; @@ -127,6 +127,7 @@ pub type JsonStream = Pin> + Se 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"; @@ -190,6 +191,7 @@ 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>; @@ -316,6 +318,7 @@ struct WorkerHandlers { subscribers: HashMap, event_metadata_injectors: HashMap, export_activation_policies: HashMap, + export_targets: HashMap, mark_sanitizers: HashMap, scope_start_sanitizers: HashMap, scope_end_sanitizers: HashMap, @@ -393,7 +396,7 @@ impl PluginContext { ); } - /// Registers this plugin's activation-time policy for remote exporters. + /// 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, @@ -407,6 +410,31 @@ impl PluginContext { ); } + /// 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, @@ -762,6 +790,7 @@ impl PluginContext { surface: surface as i32, priority, break_chain, + export_target: None, }); } } @@ -1844,6 +1873,9 @@ impl WorkerService { 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 => { @@ -1923,6 +1955,17 @@ impl WorkerService { 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, @@ -2182,6 +2225,18 @@ impl WorkerService { }) } + 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 @@ -2810,6 +2865,7 @@ fn all_surfaces() -> 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 fb1a43833..0c6361f5d 100644 --- a/crates/worker/tests/worker_sdk_tests.rs +++ b/crates/worker/tests/worker_sdk_tests.rs @@ -20,7 +20,8 @@ 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, ExportActivationDecision, Json, + 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, @@ -132,6 +133,11 @@ async fn worker_service_enforces_auth_and_reports_registrations() { .supported_surfaces .contains(&(RegistrationSurface::ExportActivationPolicy as i32)) ); + assert!( + handshake + .supported_surfaces + .contains(&(RegistrationSurface::ExportTarget as i32)) + ); let bad_health = client .health(Request::new(HealthRequest { @@ -218,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(), 23); + assert_eq!(registrations.len(), 24); for local_name in [ "llm-sanitize-request", "llm-sanitize-response", @@ -552,6 +558,8 @@ async fn worker_service_invokes_every_registration_surface() { 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, @@ -1852,6 +1860,18 @@ impl WorkerPlugin for SurfacePlugin { }, ) }); + 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 { @@ -2720,7 +2740,7 @@ fn export_activation_invoke() -> InvokeRequest { nemo_relay_worker_proto::v1::invoke_request::Payload::ExportActivation( json_envelope( "nemo.relay.ExportActivationRequest@1", - &json!({"target_kind": "otlp_trace", "config": {"enabled": true}}), + &json!({"target_kind": "nemo_relay.otlp.trace", "config": {"enabled": true}}), ) .expect("encode export activation request"), ), @@ -2728,6 +2748,24 @@ fn export_activation_invoke() -> InvokeRequest { } } +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/grpc-worker/python/about.mdx b/docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx index 851d72e7b..52c44e947 100644 --- a/docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx +++ b/docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx @@ -220,34 +220,47 @@ 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 {} - if request.target_kind in { - ExportActivationTargetKind.OTLP_TRACE, - ExportActivationTargetKind.OTLP_LOG, - ExportActivationTargetKind.OTLP_METRIC, - }: - return ( - ExportActivationDecision.ALLOW - if config.get("telemetry_allowed") is True - else ExportActivationDecision.DENY - ) - return ExportActivationDecision.ALLOW + 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. +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 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 114481b29..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 @@ -141,29 +141,45 @@ Relay compatibility range at `>=0.8.0,<1.0`, and register one callback: ```rust use nemo_relay_plugin::{ - ExportActivationDecision, ExportActivationTargetKind, Json, PluginContext, Result, + ExportActivationDecision, ExportActivationPolicyConfig, ExportActivationTargetKind, + ExportTargetRegistration, PluginContext, Result, }; -fn register_policy(context: &mut PluginContext<'_>) -> Result<()> { +fn register_geo_gated_otlp(context: &mut PluginContext<'_>) -> Result<()> { context.register_export_activation_policy(|request| async move { - let allowed = match request.target_kind { - ExportActivationTargetKind::OtlpTrace - | ExportActivationTargetKind::OtlpLog - | ExportActivationTargetKind::OtlpMetric => { - request.config.get("telemetry_allowed") == Some(&Json::Bool(true)) - } - _ => true, - }; + 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. +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 b052ed70f..04d0e76d7 100644 --- a/docs/configure-plugins/observability/configuration.mdx +++ b/docs/configure-plugins/observability/configuration.mdx @@ -91,9 +91,10 @@ share an authority. ## Export Activation Policies -Version 4 remote targets accept an optional `activation_policy`. The -`provider` is the dynamic-plugin manifest ID of a native or worker plugin that -declares the `export_activation_policy` capability and registers one callback: +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] 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 d47f0895a..75929a5d0 100644 --- a/examples/python-grpc-worker-plugin/README.md +++ b/examples/python-grpc-worker-plugin/README.md @@ -9,7 +9,7 @@ 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. It also registers an export-activation -policy that allows a remote target only when its opaque policy config contains +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 diff --git a/examples/rust-native-plugin/README.md b/examples/rust-native-plugin/README.md index eb777cc4e..65fa9b039 100644 --- a/examples/rust-native-plugin/README.md +++ b/examples/rust-native-plugin/README.md @@ -100,7 +100,7 @@ nemo-relay --bind 127.0.0.1:4040 The example registers the following runtime behavior: -- An export-activation policy that allows a remote target only when its opaque +- 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. 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 1142a8a85..2a540ce3b 100644 --- a/go/nemo_relay/observability_plugin.go +++ b/go/nemo_relay/observability_plugin.go @@ -8,7 +8,7 @@ import "encoding/json" // ObservabilityPluginKind is the top-level plugin kind used by the core observability component. const ObservabilityPluginKind = "observability" -// ExportActivationPolicyConfig attaches one dynamic policy provider to a remote exporter target. +// ExportActivationPolicyConfig attaches one policy provider to an export target. type ExportActivationPolicyConfig struct { Provider string `json:"provider"` TimeoutMillis uint64 `json:"timeout_millis,omitempty"` @@ -107,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() { @@ -156,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. diff --git a/go/nemo_relay/observability_plugin_test.go b/go/nemo_relay/observability_plugin_test.go index ac48dda20..0fcdc2eee 100644 --- a/go/nemo_relay/observability_plugin_test.go +++ b/go/nemo_relay/observability_plugin_test.go @@ -56,6 +56,7 @@ func TestObservabilityConfigHelpers(t *testing.T) { 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/" @@ -114,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) } @@ -124,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" { @@ -315,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) @@ -431,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 } @@ -442,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/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 04c501b49..6a7365c88 100644 --- a/python/nemo_relay/observability.py +++ b/python/nemo_relay/observability.py @@ -54,7 +54,7 @@ def to_dict(self) -> JsonObject: @dataclass(slots=True) class ExportActivationPolicyConfig: - """Activation-time policy attached to one remote exporter target.""" + """Activation-time policy attached to one export target.""" provider: str timeout_millis: int = 30000 @@ -123,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( @@ -131,6 +132,7 @@ def to_dict(self) -> JsonObject: "output_directory": self.output_directory, "filename": self.filename, "mode": self.mode, + "activation_policy": self.activation_policy, } ) @@ -215,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.""" @@ -228,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") diff --git a/python/nemo_relay/observability.pyi b/python/nemo_relay/observability.pyi index 9ad4cf909..447272188 100644 --- a/python/nemo_relay/observability.pyi +++ b/python/nemo_relay/observability.pyi @@ -47,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 @@ -84,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) 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 e9deb1a88..02c851409 100644 --- a/python/plugin/README.md +++ b/python/plugin/README.md @@ -35,7 +35,7 @@ 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 remote exporters at activation**: Register one typed export policy +- **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. diff --git a/python/plugin/src/nemo_relay_plugin/__init__.py b/python/plugin/src/nemo_relay_plugin/__init__.py index 02fb5c230..ba34750d7 100644 --- a/python/plugin/src/nemo_relay_plugin/__init__.py +++ b/python/plugin/src/nemo_relay_plugin/__init__.py @@ -45,9 +45,11 @@ ToolExecutionInterceptOutcome: Canonical tool execution-intercept result. DiagnosticLevel: Severity of a configuration diagnostic. ConfigDiagnostic: Structured configuration warning or error. - ExportActivationTargetKind: Remote exporter kind presented to a policy. - ExportActivationRequest: Secret-free activation request for one remote target. + 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. @@ -56,7 +58,8 @@ Public callback aliases: SubscriberCallback: Event subscriber callback. EventMetadataInjectorCallback: Event metadata injector callback. - ExportActivationPolicyCallback: Activation-time remote-export policy 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. @@ -92,8 +95,11 @@ EventSanitizeFields, ExportActivationDecision, ExportActivationPolicyCallback, + ExportActivationPolicyConfig, ExportActivationRequest, ExportActivationTargetKind, + ExportTargetActivationCallback, + ExportTargetRegistration, Json, LlmCodecIdentity, LlmConditionalCallback, @@ -151,8 +157,11 @@ "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 b9845a318..776e3962d 100644 --- a/python/plugin/src/nemo_relay_plugin/_api.py +++ b/python/plugin/src/nemo_relay_plugin/_api.py @@ -27,8 +27,8 @@ LlmRequestInterceptOutcome: Canonical LLM request-intercept result. DiagnosticLevel: Severity of a configuration diagnostic. ConfigDiagnostic: Structured configuration warning or error. - ExportActivationTargetKind: Remote exporter kind presented to a policy. - ExportActivationRequest: Secret-free activation request for one remote target. + 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. @@ -44,7 +44,7 @@ Public callback aliases used in registration annotations: SubscriberCallback: Event subscriber callback. EventMetadataInjectorCallback: Event metadata injector callback. - ExportActivationPolicyCallback: Activation-time remote-export policy callback. + ExportActivationPolicyCallback: Activation-time export policy callback. ToolSanitizeCallback: Tool request or response sanitizer callback. ToolConditionalCallback: Tool execution guardrail callback. ToolRequestCallback: Tool request intercept callback. @@ -944,15 +944,40 @@ def register(self, ctx: PluginContext, config: Json) -> None | Awaitable[None]: EventMetadataInjectorCallback: TypeAlias = Callable[[Event], dict[str, Json] | Awaitable[dict[str, Json]]] -class ExportActivationTargetKind(str, Enum): - """Remote exporter target presented to an activation policy.""" +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) + - OTLP_TRACE = "otlp_trace" - OTLP_LOG = "otlp_log" - OTLP_METRIC = "otlp_metric" - ATOF_STREAM = "atof_stream" - ATIF_HTTP = "atif_http" - ATIF_S3 = "atif_s3" +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): @@ -970,10 +995,29 @@ class ExportActivationRequest: 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], @@ -1009,6 +1053,7 @@ class _Handlers: 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] @@ -1031,6 +1076,7 @@ def empty(cls) -> _Handlers: subscribers={}, event_metadata_injectors={}, export_activation_policies={}, + export_targets={}, mark_sanitizers={}, scope_start_sanitizers={}, scope_end_sanitizers={}, @@ -1137,11 +1183,28 @@ def register_event_metadata_injector( self._handlers.event_metadata_injectors[name] = callback def register_export_activation_policy(self, callback: ExportActivationPolicyCallback) -> None: - """Register this plugin's activation-time policy for remote exporters.""" + """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, @@ -1449,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 @@ -1461,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, ) ) @@ -2299,6 +2371,9 @@ async def _invoke_result(self, request: Any) -> Any: 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 = { @@ -2484,6 +2559,7 @@ def _all_surfaces() -> list[int]: 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 1ea3ece57..0df03409a 100644 --- a/python/tests/plugin/test_public_api_docstrings.py +++ b/python/tests/plugin/test_public_api_docstrings.py @@ -24,6 +24,7 @@ "EventMetadataInjectorCallback", "EventSanitizeCallback", "ExportActivationPolicyCallback", + "ExportTargetActivationCallback", "Json", "LlmRequest", "SubscriberCallback", diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py index e46d91da3..1693218c5 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -28,6 +28,9 @@ DataSchema, DiagnosticLevel, ExportActivationDecision, + ExportActivationPolicyConfig, + ExportActivationTargetKind, + ExportTargetRegistration, Json, LlmOptimizationContribution, LlmRequestInterceptOutcome, @@ -603,6 +606,17 @@ 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) @@ -714,6 +728,7 @@ async def test_health_handshake_validate_register_and_all_surfaces(service: _Wor ("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), @@ -1657,13 +1672,23 @@ async def test_unary_invoke_success_paths(service: _WorkerService, host_stub: Re pb.EXPORT_ACTIVATION_POLICY, export_activation=_json_envelope( EXPORT_ACTIVATION_REQUEST_SCHEMA, - {"target_kind": "otlp_trace", "config": {"enabled": True}}, + {"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) @@ -3232,6 +3257,7 @@ def _all_expected_surfaces() -> list[int]: 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 be3346cb2..8780736ee 100644 --- a/python/tests/test_observability_plugin.py +++ b/python/tests/test_observability_plugin.py @@ -116,6 +116,8 @@ def test_remote_targets_serialize_export_activation_policy(self): ] == 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 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)