1use std::collections::{HashMap, HashSet};
20use std::future::Future;
21use std::path::PathBuf;
22use std::pin::Pin;
23use std::sync::{Arc, Mutex};
24#[cfg(any(feature = "otel", feature = "openinference", feature = "object-store"))]
25use std::time::Duration;
26
27use serde::{Deserialize, Serialize};
28use serde_json::{Map, Value as Json};
29use uuid::Uuid;
30
31use crate::api::event::{Event, ScopeCategory};
32use crate::api::runtime::{EventSubscriberFn, current_scope_stack};
33use crate::api::scope::ScopeType;
34use crate::api::subscriber::{
35 scope_deregister_subscriber, try_scope_deregister_subscriber, try_scope_register_subscriber,
36};
37use crate::error::FlowError;
38use crate::observability::atif::{AtifAgentInfo, AtifExporter};
39use crate::observability::atof::{
40 AtofEndpointConfig as CoreAtofEndpointConfig, AtofEndpointTransport, AtofExporter,
41 AtofExporterConfig as CoreAtofExporterConfig, AtofExporterMode,
42};
43#[cfg(feature = "openinference")]
44use crate::observability::openinference::{
45 OpenInferenceConfig as CoreOpenInferenceConfig, OpenInferenceSubscriber,
46 OtlpTransport as OpenInferenceTransport,
47};
48#[cfg(feature = "otel")]
49use crate::observability::otel::{
50 OpenTelemetryConfig as CoreOpenTelemetryConfig, OpenTelemetrySubscriber,
51};
52use crate::plugin::{
53 ConfigDiagnostic, ConfigPolicy, DiagnosticLevel, Plugin, PluginComponentSpec, PluginError,
54 PluginRegistration, PluginRegistrationContext, Result as PluginResult, UnsupportedBehavior,
55 deregister_plugin, register_plugin,
56};
57
58pub const OBSERVABILITY_PLUGIN_KIND: &str = "observability";
60
61#[derive(Debug, Clone)]
67pub struct ComponentSpec {
68 pub enabled: bool,
70 pub config: ObservabilityConfig,
72}
73
74impl ComponentSpec {
75 pub fn new(config: ObservabilityConfig) -> Self {
80 Self {
81 enabled: true,
82 config,
83 }
84 }
85}
86
87impl From<ComponentSpec> for PluginComponentSpec {
88 fn from(value: ComponentSpec) -> Self {
89 let Json::Object(config) = serde_json::to_value(value.config)
90 .expect("observability config should serialize to object")
91 else {
92 unreachable!("observability config must serialize to object");
93 };
94
95 PluginComponentSpec {
96 kind: OBSERVABILITY_PLUGIN_KIND.to_string(),
97 enabled: value.enabled,
98 config,
99 }
100 }
101}
102
103#[derive(Debug, Clone, Serialize, Deserialize)]
109#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
110pub struct ObservabilityConfig {
111 #[serde(default = "default_observability_config_version")]
113 pub version: u32,
114 #[serde(default, skip_serializing_if = "Option::is_none")]
116 pub atof: Option<AtofSectionConfig>,
117 #[serde(default, skip_serializing_if = "Option::is_none")]
119 pub atif: Option<AtifSectionConfig>,
120 #[serde(default, skip_serializing_if = "Option::is_none")]
122 pub opentelemetry: Option<OtlpSectionConfig>,
123 #[serde(default, skip_serializing_if = "Option::is_none")]
125 pub openinference: Option<OtlpSectionConfig>,
126 #[serde(default)]
128 pub policy: ConfigPolicy,
129}
130
131impl Default for ObservabilityConfig {
132 fn default() -> Self {
133 Self {
134 version: default_observability_config_version(),
135 atof: None,
136 atif: None,
137 opentelemetry: None,
138 openinference: None,
139 policy: ConfigPolicy::default(),
140 }
141 }
142}
143
144#[derive(Debug, Clone, Serialize, Deserialize)]
151#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
152pub struct AtofSectionConfig {
153 #[serde(default)]
155 pub enabled: bool,
156 #[serde(default, skip_serializing_if = "Option::is_none")]
158 pub output_directory: Option<PathBuf>,
159 #[serde(default, skip_serializing_if = "Option::is_none")]
161 pub filename: Option<String>,
162 #[serde(default = "default_atof_mode")]
164 #[cfg_attr(feature = "schema", schemars(schema_with = "atof_mode_schema"))]
165 pub mode: String,
166 #[serde(default, skip_serializing_if = "Vec::is_empty")]
168 pub endpoints: Vec<AtofEndpointSectionConfig>,
169}
170
171impl Default for AtofSectionConfig {
172 fn default() -> Self {
173 Self {
174 enabled: false,
175 output_directory: None,
176 filename: None,
177 mode: default_atof_mode(),
178 endpoints: Vec::new(),
179 }
180 }
181}
182
183#[derive(Debug, Clone, Serialize, Deserialize)]
185#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
186pub struct AtofEndpointSectionConfig {
187 pub url: String,
189 #[serde(default = "default_atof_endpoint_transport")]
191 #[cfg_attr(
192 feature = "schema",
193 schemars(schema_with = "atof_endpoint_transport_schema")
194 )]
195 pub transport: String,
196 #[serde(default)]
198 pub headers: HashMap<String, String>,
199 #[serde(default = "default_timeout_millis")]
201 pub timeout_millis: u64,
202}
203
204#[derive(Debug, Clone, Serialize, Deserialize)]
211#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
212pub struct AtifSectionConfig {
213 #[serde(default)]
215 pub enabled: bool,
216 #[serde(default = "default_agent_name")]
218 pub agent_name: String,
219 #[serde(default = "default_agent_version")]
221 pub agent_version: String,
222 #[serde(default = "default_model_name")]
224 pub model_name: String,
225 #[serde(default, skip_serializing_if = "Option::is_none")]
227 pub tool_definitions: Option<Vec<Json>>,
228 #[serde(default, skip_serializing_if = "Option::is_none")]
230 pub extra: Option<Json>,
231 #[serde(default, skip_serializing_if = "Option::is_none")]
235 pub output_directory: Option<PathBuf>,
236 #[serde(default = "default_atif_filename_template")]
241 pub filename_template: String,
242 #[serde(default, skip_serializing_if = "Vec::is_empty")]
251 pub storage: Vec<AtifStorageConfig>,
252}
253
254impl Default for AtifSectionConfig {
255 fn default() -> Self {
256 Self {
257 enabled: false,
258 agent_name: default_agent_name(),
259 agent_version: default_agent_version(),
260 model_name: default_model_name(),
261 tool_definitions: None,
262 extra: None,
263 output_directory: None,
264 filename_template: default_atif_filename_template(),
265 storage: Vec::new(),
266 }
267 }
268}
269
270#[derive(Debug, Clone, Serialize, Deserialize)]
278#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
279#[serde(tag = "type", rename_all = "snake_case")]
280pub enum AtifStorageConfig {
281 Http(HttpStorageConfig),
283 S3(S3StorageConfig),
295}
296
297#[derive(Debug, Clone, Serialize, Deserialize)]
306#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
307pub struct S3StorageConfig {
308 pub bucket: String,
310 #[serde(default, skip_serializing_if = "Option::is_none")]
313 pub key_prefix: Option<String>,
314 #[serde(default, skip_serializing_if = "Option::is_none")]
316 pub access_key_id: Option<String>,
317 #[serde(default, skip_serializing_if = "Option::is_none")]
321 pub secret_access_key_var: Option<String>,
322 #[serde(default, skip_serializing_if = "Option::is_none")]
326 pub session_token_var: Option<String>,
327 #[serde(default, skip_serializing_if = "Option::is_none")]
329 pub region: Option<String>,
330 #[serde(default, skip_serializing_if = "Option::is_none")]
333 pub endpoint_url: Option<String>,
334 #[serde(default, skip_serializing_if = "Option::is_none")]
336 pub allow_http: Option<bool>,
337}
338
339#[derive(Debug, Clone, Serialize, Deserialize)]
346#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
347pub struct HttpStorageConfig {
348 pub endpoint: String,
350 #[serde(default, skip_serializing_if = "HashMap::is_empty")]
352 pub headers: HashMap<String, String>,
353 #[serde(default, skip_serializing_if = "HashMap::is_empty")]
355 pub header_env: HashMap<String, String>,
356 #[serde(default = "default_timeout_millis")]
358 pub timeout_millis: u64,
359}
360
361#[derive(Debug, Clone, Serialize, Deserialize)]
367#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
368pub struct OtlpSectionConfig {
369 #[serde(default)]
371 pub enabled: bool,
372 #[serde(default = "default_otlp_transport")]
374 #[cfg_attr(feature = "schema", schemars(schema_with = "otlp_transport_schema"))]
375 pub transport: String,
376 #[serde(default, skip_serializing_if = "Option::is_none")]
378 pub endpoint: Option<String>,
379 #[serde(default)]
381 pub headers: HashMap<String, String>,
382 #[serde(default)]
384 pub resource_attributes: HashMap<String, String>,
385 #[serde(default = "default_service_name")]
387 pub service_name: String,
388 #[serde(default, skip_serializing_if = "Option::is_none")]
390 pub service_namespace: Option<String>,
391 #[serde(default, skip_serializing_if = "Option::is_none")]
393 pub service_version: Option<String>,
394 #[serde(default, skip_serializing_if = "Option::is_none")]
396 pub instrumentation_scope: Option<String>,
397 #[serde(default = "default_timeout_millis")]
399 pub timeout_millis: u64,
400}
401
402impl Default for OtlpSectionConfig {
403 fn default() -> Self {
404 Self {
405 enabled: false,
406 transport: default_otlp_transport(),
407 endpoint: None,
408 headers: HashMap::new(),
409 resource_attributes: HashMap::new(),
410 service_name: default_service_name(),
411 service_namespace: None,
412 service_version: None,
413 instrumentation_scope: None,
414 timeout_millis: default_timeout_millis(),
415 }
416 }
417}
418
419crate::editor_config! {
420 impl ObservabilityConfig {
421 atof => {
422 label: "ATOF",
423 kind: Section,
424 optional: true,
425 nested: AtofSectionConfig,
426 default: AtofSectionConfig,
427 },
428 atif => {
429 label: "ATIF",
430 kind: Section,
431 optional: true,
432 nested: AtifSectionConfig,
433 default: AtifSectionConfig,
434 },
435 opentelemetry => {
436 label: "OpenTelemetry",
437 kind: Section,
438 optional: true,
439 nested: OtlpSectionConfig,
440 default: OtlpSectionConfig,
441 },
442 openinference => {
443 label: "OpenInference",
444 kind: Section,
445 optional: true,
446 nested: OtlpSectionConfig,
447 default: OtlpSectionConfig,
448 },
449 policy => {
450 label: "policy",
451 kind: Section,
452 nested: ConfigPolicy,
453 default: ConfigPolicy,
454 },
455 }
456}
457
458crate::editor_config! {
459 impl AtofSectionConfig {
460 enabled => { label: "enabled", kind: Boolean },
461 output_directory => { label: "output_directory", kind: String, optional: true },
462 filename => { label: "filename", kind: String, optional: true },
463 mode => { label: "mode", kind: Enum, values: ["append", "overwrite"] },
464 endpoints => { label: "endpoints", kind: Json, optional: true },
465 }
466}
467
468crate::editor_config! {
469 impl AtifSectionConfig {
470 enabled => { label: "enabled", kind: Boolean },
471 agent_name => { label: "agent_name", kind: String },
472 agent_version => { label: "agent_version", kind: String },
473 model_name => { label: "model_name", kind: String },
474 tool_definitions => { label: "tool_definitions", kind: Json, optional: true },
475 extra => { label: "extra", kind: Json, optional: true },
476 output_directory => { label: "output_directory", kind: String, optional: true },
477 filename_template => { label: "filename_template", kind: String },
478 storage => { label: "storage", kind: Json, optional: true },
479 }
480}
481
482crate::editor_config! {
483 impl OtlpSectionConfig {
484 enabled => { label: "enabled", kind: Boolean },
485 transport => { label: "transport", kind: Enum, values: ["http_binary", "grpc"] },
486 endpoint => { label: "endpoint", kind: String, optional: true },
487 headers => { label: "headers", kind: StringMap },
488 resource_attributes => { label: "resource_attributes", kind: StringMap },
489 service_name => { label: "service_name", kind: String },
490 service_namespace => { label: "service_namespace", kind: String, optional: true },
491 service_version => { label: "service_version", kind: String, optional: true },
492 instrumentation_scope => { label: "instrumentation_scope", kind: String, optional: true },
493 timeout_millis => { label: "timeout_millis", kind: Integer },
494 }
495}
496
497struct ObservabilityPlugin;
498
499impl Plugin for ObservabilityPlugin {
500 fn plugin_kind(&self) -> &str {
501 OBSERVABILITY_PLUGIN_KIND
502 }
503
504 fn allows_multiple_components(&self) -> bool {
505 false
506 }
507
508 fn validate(&self, plugin_config: &Map<String, Json>) -> Vec<ConfigDiagnostic> {
509 validate_observability_plugin_config(plugin_config)
510 }
511
512 fn register<'a>(
513 &'a self,
514 plugin_config: &Map<String, Json>,
515 ctx: &'a mut PluginRegistrationContext,
516 ) -> Pin<Box<dyn Future<Output = PluginResult<()>> + Send + 'a>> {
517 let plugin_config = plugin_config.clone();
518 Box::pin(async move {
519 let config = parse_observability_config(&plugin_config)?;
520 register_observability(config, ctx)
521 })
522 }
523}
524
525pub fn register_observability_component() -> PluginResult<()> {
531 match register_plugin(Arc::new(ObservabilityPlugin)) {
532 Ok(()) => Ok(()),
533 Err(PluginError::RegistrationFailed(message)) if message.contains("already registered") => {
534 Ok(())
535 }
536 Err(err) => Err(err),
537 }
538}
539
540pub fn deregister_observability_component() -> bool {
546 deregister_plugin(OBSERVABILITY_PLUGIN_KIND)
547}
548
549#[cfg(feature = "schema")]
551pub fn observability_config_schema() -> serde_json::Value {
552 serde_json::to_value(schemars::schema_for!(ObservabilityConfig))
553 .expect("observability config schema should serialize")
554}
555
556#[cfg(feature = "schema")]
557fn atof_mode_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::schema::Schema {
558 string_enum_schema(generator, &["append", "overwrite"], Some("append"))
559}
560
561#[cfg(feature = "schema")]
562fn atof_endpoint_transport_schema(
563 generator: &mut schemars::r#gen::SchemaGenerator,
564) -> schemars::schema::Schema {
565 string_enum_schema(
566 generator,
567 &["http_post", "websocket", "ndjson"],
568 Some("http_post"),
569 )
570}
571
572#[cfg(feature = "schema")]
573fn otlp_transport_schema(
574 generator: &mut schemars::r#gen::SchemaGenerator,
575) -> schemars::schema::Schema {
576 string_enum_schema(generator, &["http_binary", "grpc"], Some("http_binary"))
577}
578
579#[cfg(feature = "schema")]
580fn string_enum_schema(
581 generator: &mut schemars::r#gen::SchemaGenerator,
582 values: &[&str],
583 default: Option<&str>,
584) -> schemars::schema::Schema {
585 let mut schema: schemars::schema::SchemaObject =
586 <String as schemars::JsonSchema>::json_schema(generator).into();
587 schema.enum_values = Some(
588 values
589 .iter()
590 .map(|value| Json::String((*value).into()))
591 .collect(),
592 );
593 if let Some(default) = default {
594 schema.metadata().default = Some(Json::String(default.into()));
595 }
596 schema.into()
597}
598
599fn register_observability(
600 config: ObservabilityConfig,
601 ctx: &mut PluginRegistrationContext,
602) -> PluginResult<()> {
603 if let Some(atof) = config.atof.filter(|section| section.enabled) {
604 register_atof_exporter(atof, ctx)?;
605 }
606 if let Some(atif) = config.atif.filter(|section| section.enabled) {
607 register_atif_dispatcher(atif, ctx)?;
608 }
609 if let Some(otel) = config.opentelemetry.filter(|section| section.enabled) {
610 register_opentelemetry(otel, ctx)?;
611 }
612 if let Some(openinference) = config.openinference.filter(|section| section.enabled) {
613 register_openinference(openinference, ctx)?;
614 }
615 Ok(())
616}
617
618fn register_atof_exporter(
619 section: AtofSectionConfig,
620 ctx: &mut PluginRegistrationContext,
621) -> PluginResult<()> {
622 let mode = AtofExporterMode::parse(§ion.mode).ok_or_else(|| {
623 PluginError::InvalidConfig("ATOF mode must be 'append' or 'overwrite'".to_string())
624 })?;
625 let mut config = CoreAtofExporterConfig::new().with_mode(mode);
626 if let Some(output_directory) = section.output_directory {
627 config = config.with_output_directory(output_directory);
628 }
629 if let Some(filename) = section.filename {
630 config = config.with_filename(filename);
631 }
632 let endpoints = section
633 .endpoints
634 .into_iter()
635 .enumerate()
636 .map(|(index, endpoint)| build_atof_endpoint_config(index, endpoint))
637 .collect::<PluginResult<Vec<_>>>()?;
638 config = config.with_endpoints(endpoints);
639
640 let exporter = Arc::new(AtofExporter::new(config).map_err(observability_registration_error)?);
641 ctx.register_subscriber("atof", exporter.subscriber())?;
642 ctx.add_registration(PluginRegistration::new(
643 "observability",
644 ctx.qualify_name("atof.shutdown"),
645 Box::new(move || {
646 exporter
647 .shutdown()
648 .map_err(observability_registration_error)
649 }),
650 ));
651 Ok(())
652}
653
654fn build_atof_endpoint_config(
655 index: usize,
656 endpoint: AtofEndpointSectionConfig,
657) -> PluginResult<CoreAtofEndpointConfig> {
658 let transport = AtofEndpointTransport::parse(&endpoint.transport).ok_or_else(|| {
659 PluginError::InvalidConfig(format!(
660 "ATOF endpoints[{index}].transport must be 'http_post', 'websocket', or 'ndjson'"
661 ))
662 })?;
663 let mut config = CoreAtofEndpointConfig::new(endpoint.url, transport)
664 .with_timeout_millis(endpoint.timeout_millis);
665 for (key, value) in endpoint.headers {
666 config = config.with_header(key, value);
667 }
668 Ok(config)
669}
670
671type AtifStorageList = Arc<Vec<Arc<AtifRemoteStorage>>>;
672
673fn register_atif_dispatcher(
674 section: AtifSectionConfig,
675 ctx: &mut PluginRegistrationContext,
676) -> PluginResult<()> {
677 if !section.filename_template.contains("{session_id}") {
678 return Err(PluginError::InvalidConfig(
679 "ATIF filename_template must contain '{session_id}'".to_string(),
680 ));
681 }
682
683 let mut storage_vec = Vec::with_capacity(section.storage.len());
684 for (index, entry) in section.storage.iter().enumerate() {
685 storage_vec.push(build_atif_storage(index, entry)?);
686 }
687 let storage: AtifStorageList = Arc::new(storage_vec);
688
689 let manager = Arc::new(Mutex::new(AtifDispatcher::new(section)));
690 let dispatcher = atif_dispatcher_subscriber(
691 Arc::clone(&manager),
692 ctx.qualify_name("atif-"),
693 Arc::clone(&storage),
694 );
695 ctx.register_subscriber("atif", dispatcher)?;
696 let shutdown_storage = Arc::clone(&storage);
697 ctx.add_registration(PluginRegistration::new(
698 "observability",
699 ctx.qualify_name("atif.shutdown"),
700 Box::new(move || {
701 let work = {
702 let mut guard = manager.lock().map_err(|err| {
703 PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}"))
704 })?;
705 guard.flush_open_agents()
706 };
707 for (scope_uuid, name) in work.scope_subscribers {
708 deregister_atif_shutdown_subscriber(&scope_uuid, &name)?;
709 }
710 for export in work.exports {
711 let write = prepare_atif_shutdown_file(&export, Arc::clone(&manager))
712 .map_err(observability_registration_error)?;
713 let agent_uuid = write.agent_uuid;
714 let targets = {
715 let guard = manager.lock().map_err(|err| {
716 PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}"))
717 })?;
718 guard.sink_targets()
719 };
720 let results = write_atif(&write, shutdown_storage.as_slice(), &targets);
721 let mut guard = manager.lock().map_err(|err| {
722 PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}"))
723 })?;
724 let _ = guard.complete_scope_write(agent_uuid, results);
725 }
726 let guard = manager.lock().map_err(|err| {
727 PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}"))
728 })?;
729 guard
730 .last_error_result()
731 .map_err(observability_registration_error)
732 }),
733 ));
734 Ok(())
735}
736
737fn deregister_atif_shutdown_subscriber(scope_uuid: &Uuid, name: &str) -> PluginResult<()> {
738 match scope_deregister_subscriber(scope_uuid, name) {
739 Ok(_) | Err(FlowError::NotFound(_)) => Ok(()),
740 Err(error) => Err(observability_registration_error(error)),
741 }
742}
743
744#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
745fn build_atif_storage(
746 index: usize,
747 config: &AtifStorageConfig,
748) -> PluginResult<Arc<AtifRemoteStorage>> {
749 AtifRemoteStorage::from_config(index, config)
750 .map(Arc::new)
751 .map_err(observability_registration_error)
752}
753
754#[cfg(not(all(feature = "object-store", not(target_arch = "wasm32"))))]
755fn build_atif_storage(
756 _index: usize,
757 _config: &AtifStorageConfig,
758) -> PluginResult<Arc<AtifRemoteStorage>> {
759 Err(PluginError::InvalidConfig(
760 "ATIF storage support is not enabled in this build".to_string(),
761 ))
762}
763
764#[cfg(feature = "otel")]
765fn register_opentelemetry(
766 section: OtlpSectionConfig,
767 ctx: &mut PluginRegistrationContext,
768) -> PluginResult<()> {
769 let subscriber = Arc::new(
770 OpenTelemetrySubscriber::new(build_otel_config(section)?)
771 .map_err(observability_registration_error)?,
772 );
773 ctx.register_subscriber("opentelemetry", subscriber.subscriber())?;
774 ctx.add_registration(PluginRegistration::new(
775 "observability",
776 ctx.qualify_name("opentelemetry.shutdown"),
777 Box::new(move || {
778 subscriber
779 .shutdown()
780 .map_err(observability_registration_error)
781 }),
782 ));
783 Ok(())
784}
785
786#[cfg(not(feature = "otel"))]
787fn register_opentelemetry(
788 _section: OtlpSectionConfig,
789 _ctx: &mut PluginRegistrationContext,
790) -> PluginResult<()> {
791 Err(PluginError::InvalidConfig(
792 "OpenTelemetry support is not enabled in this build".to_string(),
793 ))
794}
795
796#[cfg(feature = "openinference")]
797fn register_openinference(
798 section: OtlpSectionConfig,
799 ctx: &mut PluginRegistrationContext,
800) -> PluginResult<()> {
801 let subscriber = Arc::new(
802 OpenInferenceSubscriber::new(build_openinference_config(section)?)
803 .map_err(observability_registration_error)?,
804 );
805 ctx.register_subscriber("openinference", subscriber.subscriber())?;
806 ctx.add_registration(PluginRegistration::new(
807 "observability",
808 ctx.qualify_name("openinference.shutdown"),
809 Box::new(move || {
810 subscriber
811 .shutdown()
812 .map_err(observability_registration_error)
813 }),
814 ));
815 Ok(())
816}
817
818#[cfg(not(feature = "openinference"))]
819fn register_openinference(
820 _section: OtlpSectionConfig,
821 _ctx: &mut PluginRegistrationContext,
822) -> PluginResult<()> {
823 Err(PluginError::InvalidConfig(
824 "OpenInference support is not enabled in this build".to_string(),
825 ))
826}
827
828struct AtifDispatcher {
829 config: AtifSectionConfig,
830 agents: HashMap<Uuid, ManagedAtifExporter>,
831 scope_owners: HashMap<Uuid, Uuid>,
832 scope_subscribers: HashMap<Uuid, String>,
833 fatal_error: Option<String>,
837 sink_errors: HashMap<SinkLabel, String>,
840}
841
842struct ManagedAtifExporter {
843 exporter: AtifExporter,
844 filename: String,
845 local_path: Option<PathBuf>,
846 observed_events: Vec<Event>,
847 observed_event_keys: HashSet<String>,
848 written: bool,
849}
850
851struct PendingAtifWrite {
852 agent_uuid: Uuid,
853 #[cfg_attr(
854 not(all(feature = "object-store", not(target_arch = "wasm32"))),
855 allow(dead_code)
856 )]
857 session_id: String,
858 #[cfg_attr(
861 not(all(feature = "object-store", not(target_arch = "wasm32"))),
862 allow(dead_code)
863 )]
864 filename: String,
865 local_path: Option<PathBuf>,
866 payload: Vec<u8>,
867}
868
869struct AtifFlushWork {
870 exports: Vec<PendingAtifExport>,
871 scope_subscribers: Vec<(Uuid, String)>,
872}
873
874struct PendingAtifExport {
875 agent_uuid: Uuid,
876 exporter: AtifExporter,
877 filename: String,
878 local_path: Option<PathBuf>,
879}
880
881#[derive(Clone, Debug, PartialEq, Eq, Hash)]
885enum SinkLabel {
886 Local,
887 Remote(usize),
888}
889
890impl AtifDispatcher {
891 fn new(config: AtifSectionConfig) -> Self {
892 Self {
893 config,
894 agents: HashMap::new(),
895 scope_owners: HashMap::new(),
896 scope_subscribers: HashMap::new(),
897 fatal_error: None,
898 sink_errors: HashMap::new(),
899 }
900 }
901
902 fn observe_global(
903 &mut self,
904 event: &Event,
905 subscriber_prefix: &str,
906 state: Arc<Mutex<Self>>,
907 storage: AtifStorageList,
908 ) -> Option<(PendingAtifWrite, Vec<SinkLabel>)> {
909 if self.fatal_error.is_some() {
910 return None;
911 }
912
913 if !is_top_level_trajectory_start(event) {
914 return self.observe_descendant_from_global(event);
915 }
916
917 if self.agents.contains_key(&event.uuid()) {
918 return None;
919 }
920
921 let session_id = event.uuid().to_string();
926 let exporter = AtifExporter::new(session_id.clone(), self.agent_info());
927 (exporter.subscriber())(event);
928 let (filename, local_path) = self.prepare_destination(&session_id);
929 self.scope_owners.insert(event.uuid(), event.uuid());
930 self.agents.insert(
931 event.uuid(),
932 ManagedAtifExporter {
933 exporter,
934 filename,
935 local_path,
936 observed_events: vec![event.clone()],
937 observed_event_keys: HashSet::from([event_observation_key(event)]),
938 written: false,
939 },
940 );
941
942 let agent_uuid = event.uuid();
943 let name = format!("{subscriber_prefix}{agent_uuid}");
944 let callback = atif_scope_subscriber(state, agent_uuid, storage);
945 if try_scope_register_subscriber(&agent_uuid, &name, callback).is_ok() {
951 self.scope_subscribers.insert(agent_uuid, name);
952 }
953 None
954 }
955
956 fn observe_descendant_from_global(
957 &mut self,
958 event: &Event,
959 ) -> Option<(PendingAtifWrite, Vec<SinkLabel>)> {
960 let owner = self.scope_owners.get(&event.uuid()).copied().or_else(|| {
961 event
962 .parent_uuid()
963 .and_then(|uuid| self.scope_owners.get(&uuid).copied())
964 })?;
965
966 if event.scope_category() == Some(ScopeCategory::Start) {
967 self.scope_owners.insert(event.uuid(), owner);
968 }
969
970 let pending_write = self.observe_scope(event, owner);
971
972 if event.scope_category() == Some(ScopeCategory::End) && event.uuid() != owner {
973 self.scope_owners.remove(&event.uuid());
974 }
975
976 pending_write
977 }
978
979 fn observe_scope(
980 &mut self,
981 event: &Event,
982 agent_uuid: Uuid,
983 ) -> Option<(PendingAtifWrite, Vec<SinkLabel>)> {
984 if self.fatal_error.is_some() {
985 return None;
986 }
987 let should_finalize =
988 event.uuid() == agent_uuid && event.scope_category() == Some(ScopeCategory::End);
989 let agent = self.agents.get_mut(&agent_uuid)?;
990 if !agent
991 .observed_event_keys
992 .insert(event_observation_key(event))
993 {
994 return None;
995 }
996 (agent.exporter.subscriber())(event);
997 agent.observed_events.push(event.clone());
998 if !should_finalize || agent.written {
999 return None;
1000 }
1001 let write = match prepare_atif_file(agent_uuid, agent) {
1002 Ok(write) => write,
1003 Err(err) => {
1004 self.fatal_error = Some(err.to_string());
1005 return None;
1006 }
1007 };
1008 let targets = self.sink_targets();
1009 Some((write, targets))
1010 }
1011
1012 fn complete_scope_write(
1013 &mut self,
1014 agent_uuid: Uuid,
1015 results: Vec<(SinkLabel, std::io::Result<()>)>,
1016 ) -> Option<(Uuid, String)> {
1017 for (label, result) in results {
1018 if let Err(err) = result {
1019 self.sink_errors.insert(label, err.to_string());
1020 }
1021 }
1022 if let Some(agent) = self.agents.get_mut(&agent_uuid) {
1023 agent.observed_events.clear();
1024 }
1025 self.agents.remove(&agent_uuid);
1026 self.scope_owners.retain(|_, owner| *owner != agent_uuid);
1027 self.scope_subscribers
1028 .remove(&agent_uuid)
1029 .map(|name| (agent_uuid, name))
1030 }
1031
1032 fn flush_open_agents(&mut self) -> AtifFlushWork {
1033 let scope_subscribers = std::mem::take(&mut self.scope_subscribers)
1037 .into_iter()
1038 .collect();
1039 let agent_uuids = self
1040 .agents
1041 .iter()
1042 .filter_map(|(agent_uuid, agent)| (!agent.written).then_some(*agent_uuid))
1043 .collect::<Vec<_>>();
1044 let mut exports = Vec::with_capacity(agent_uuids.len());
1045 for agent_uuid in agent_uuids {
1046 if let Some(agent) = self.agents.get_mut(&agent_uuid) {
1047 agent.written = true;
1048 exports.push(PendingAtifExport {
1049 agent_uuid,
1050 exporter: agent.exporter.clone(),
1051 filename: agent.filename.clone(),
1052 local_path: agent.local_path.clone(),
1053 });
1054 }
1055 }
1056 AtifFlushWork {
1057 exports,
1058 scope_subscribers,
1059 }
1060 }
1061
1062 fn observed_events(&self, agent_uuid: Uuid) -> Vec<Event> {
1063 self.agents
1064 .get(&agent_uuid)
1065 .map(|agent| agent.observed_events.clone())
1066 .unwrap_or_default()
1067 }
1068
1069 fn last_error_result(&self) -> std::io::Result<()> {
1070 if let Some(message) = &self.fatal_error {
1071 return Err(std::io::Error::other(message.clone()));
1072 }
1073 Ok(())
1074 }
1075
1076 fn agent_info(&self) -> AtifAgentInfo {
1077 AtifAgentInfo {
1078 name: self.config.agent_name.clone(),
1079 version: self.config.agent_version.clone(),
1080 model_name: Some(self.config.model_name.clone()),
1081 tool_definitions: self.config.tool_definitions.clone(),
1082 extra: self.config.extra.clone(),
1083 }
1084 }
1085
1086 fn prepare_destination(&self, session_id: &str) -> (String, Option<PathBuf>) {
1087 let filename = self
1088 .config
1089 .filename_template
1090 .replace("{session_id}", session_id);
1091 if !self.config.storage.is_empty() {
1092 return (filename, None);
1093 }
1094 let directory = self
1095 .config
1096 .output_directory
1097 .clone()
1098 .unwrap_or_else(default_output_directory);
1099 let path = directory.join(&filename);
1100 (filename, Some(path))
1101 }
1102
1103 fn sink_targets(&self) -> Vec<SinkLabel> {
1104 if self.config.storage.is_empty() {
1105 if self.sink_errors.contains_key(&SinkLabel::Local) {
1106 Vec::new()
1107 } else {
1108 vec![SinkLabel::Local]
1109 }
1110 } else {
1111 (0..self.config.storage.len())
1112 .map(SinkLabel::Remote)
1113 .filter(|label| !self.sink_errors.contains_key(label))
1114 .collect()
1115 }
1116 }
1117}
1118
1119fn atif_dispatcher_subscriber(
1120 manager: Arc<Mutex<AtifDispatcher>>,
1121 subscriber_prefix: String,
1122 storage: AtifStorageList,
1123) -> EventSubscriberFn {
1124 Arc::new(move |event: &Event| {
1125 let pending = {
1126 let Ok(mut guard) = manager.lock() else {
1127 return;
1128 };
1129 guard.observe_global(
1130 event,
1131 &subscriber_prefix,
1132 Arc::clone(&manager),
1133 Arc::clone(&storage),
1134 )
1135 };
1136 let Some((write, targets)) = pending else {
1137 return;
1138 };
1139 let results = write_atif(&write, storage.as_slice(), &targets);
1140 let scope_subscriber = {
1141 let Ok(mut guard) = manager.lock() else {
1142 return;
1143 };
1144 guard.complete_scope_write(write.agent_uuid, results)
1145 };
1146 if let Some((scope_uuid, name)) = scope_subscriber {
1147 let _ = try_scope_deregister_subscriber(&scope_uuid, &name);
1148 }
1149 })
1150}
1151
1152fn atif_scope_subscriber(
1153 manager: Arc<Mutex<AtifDispatcher>>,
1154 agent_uuid: Uuid,
1155 storage: AtifStorageList,
1156) -> EventSubscriberFn {
1157 Arc::new(move |event: &Event| {
1158 let pending = {
1159 let Ok(mut guard) = manager.lock() else {
1160 return;
1161 };
1162 guard.observe_scope(event, agent_uuid)
1163 };
1164 let Some((write, targets)) = pending else {
1165 return;
1166 };
1167 let results = write_atif(&write, storage.as_slice(), &targets);
1168 let scope_subscriber = {
1169 let Ok(mut guard) = manager.lock() else {
1170 return;
1171 };
1172 guard.complete_scope_write(write.agent_uuid, results)
1173 };
1174 if let Some((scope_uuid, name)) = scope_subscriber {
1175 let _ = try_scope_deregister_subscriber(&scope_uuid, &name);
1176 }
1177 })
1178}
1179
1180fn prepare_atif_file(
1181 agent_uuid: Uuid,
1182 agent: &mut ManagedAtifExporter,
1183) -> std::io::Result<PendingAtifWrite> {
1184 let trajectory = agent
1185 .exporter
1186 .try_export()
1187 .map_err(|error| std::io::Error::other(error.to_string()))?;
1188 let observed_events = agent.observed_events.clone();
1189 agent.written = true;
1190 prepare_atif_payload(
1191 agent_uuid,
1192 agent.filename.clone(),
1193 agent.local_path.clone(),
1194 trajectory,
1195 observed_events,
1196 )
1197}
1198
1199fn prepare_atif_shutdown_file(
1200 export: &PendingAtifExport,
1201 manager: Arc<Mutex<AtifDispatcher>>,
1202) -> std::io::Result<PendingAtifWrite> {
1203 let trajectory = export
1204 .exporter
1205 .try_export()
1206 .map_err(|error| std::io::Error::other(error.to_string()))?;
1207 let observed_events = {
1208 let guard = manager.lock().map_err(|err| {
1209 std::io::Error::other(format!("ATIF dispatcher lock poisoned: {err}"))
1210 })?;
1211 guard.observed_events(export.agent_uuid)
1212 };
1213 prepare_atif_payload(
1214 export.agent_uuid,
1215 export.filename.clone(),
1216 export.local_path.clone(),
1217 trajectory,
1218 observed_events,
1219 )
1220}
1221
1222fn prepare_atif_payload(
1223 agent_uuid: Uuid,
1224 filename: String,
1225 local_path: Option<PathBuf>,
1226 trajectory: crate::observability::atif::AtifTrajectory,
1227 observed_events: Vec<Event>,
1228) -> std::io::Result<PendingAtifWrite> {
1229 let mut value = serde_json::to_value(trajectory)?;
1230 if let Some(object) = value.as_object_mut() {
1231 object.insert(
1232 "extra".to_string(),
1233 serde_json::json!({
1234 "observed_events": observed_events,
1235 }),
1236 );
1237 }
1238 let payload = serde_json::to_vec_pretty(&value)?;
1239 Ok(PendingAtifWrite {
1240 agent_uuid,
1241 session_id: agent_uuid.to_string(),
1242 filename,
1243 local_path,
1244 payload,
1245 })
1246}
1247
1248fn write_atif(
1249 write: &PendingAtifWrite,
1250 storage: &[Arc<AtifRemoteStorage>],
1251 targets: &[SinkLabel],
1252) -> Vec<(SinkLabel, std::io::Result<()>)> {
1253 targets
1254 .iter()
1255 .map(|label| {
1256 let result = match label {
1257 SinkLabel::Local => match &write.local_path {
1258 Some(path) => write_atif_local(path, &write.payload),
1259 None => Err(std::io::Error::other(
1260 "ATIF local destination has no output path",
1261 )),
1262 },
1263 SinkLabel::Remote(index) => write_atif_remote(storage, *index, write),
1264 };
1265 (label.clone(), result)
1266 })
1267 .collect()
1268}
1269
1270fn write_atif_local(path: &PathBuf, payload: &[u8]) -> std::io::Result<()> {
1271 if let Some(parent) = path.parent() {
1272 std::fs::create_dir_all(parent)?;
1273 }
1274 std::fs::write(path, payload)
1275}
1276
1277#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
1278fn write_atif_remote(
1279 storage: &[Arc<AtifRemoteStorage>],
1280 index: usize,
1281 write: &PendingAtifWrite,
1282) -> std::io::Result<()> {
1283 let sink = storage
1284 .get(index)
1285 .ok_or_else(|| std::io::Error::other(format!("ATIF storage[{index}] is not registered")))?;
1286 sink.put(&write.filename, &write.session_id, &write.payload)
1287}
1288
1289#[cfg(not(all(feature = "object-store", not(target_arch = "wasm32"))))]
1290fn write_atif_remote(
1291 _storage: &[Arc<AtifRemoteStorage>],
1292 _index: usize,
1293 _write: &PendingAtifWrite,
1294) -> std::io::Result<()> {
1295 Err(std::io::Error::other(
1296 "ATIF storage support is not enabled in this build",
1297 ))
1298}
1299
1300fn event_observation_key(event: &Event) -> String {
1301 format!(
1302 "{}:{}:{:?}",
1303 event.kind(),
1304 event.uuid(),
1305 event.scope_category()
1306 )
1307}
1308
1309fn is_top_level_trajectory_start(event: &Event) -> bool {
1310 if event.scope_category() != Some(ScopeCategory::Start) {
1311 return false;
1312 }
1313 if event.scope_type() != Some(ScopeType::Agent) {
1314 return false;
1315 }
1316 let Some(parent_uuid) = event.parent_uuid() else {
1317 return false;
1318 };
1319 current_scope_stack()
1320 .read()
1321 .map(|stack| stack.root_uuid() == parent_uuid)
1322 .unwrap_or(false)
1323}
1324
1325#[cfg(feature = "otel")]
1326fn build_otel_config(section: OtlpSectionConfig) -> PluginResult<CoreOpenTelemetryConfig> {
1327 let mut config = match section.transport.as_str() {
1328 "http_binary" => CoreOpenTelemetryConfig::http_binary(section.service_name),
1329 "grpc" => CoreOpenTelemetryConfig::grpc(section.service_name),
1330 other => {
1331 return Err(PluginError::InvalidConfig(format!(
1332 "OpenTelemetry transport must be 'http_binary' or 'grpc', got {other:?}"
1333 )));
1334 }
1335 }
1336 .with_timeout(Duration::from_millis(section.timeout_millis));
1337
1338 if let Some(endpoint) = section.endpoint {
1339 config = config.with_endpoint(endpoint);
1340 }
1341 if let Some(namespace) = section.service_namespace {
1342 config = config.with_service_namespace(namespace);
1343 }
1344 if let Some(version) = section.service_version {
1345 config = config.with_service_version(version);
1346 }
1347 if let Some(scope) = section.instrumentation_scope {
1348 config = config.with_instrumentation_scope(scope);
1349 }
1350 for (key, value) in section.headers {
1351 config = config.with_header(key, value);
1352 }
1353 for (key, value) in section.resource_attributes {
1354 config = config.with_resource_attribute(key, value);
1355 }
1356 Ok(config)
1357}
1358
1359#[cfg(feature = "openinference")]
1360fn build_openinference_config(section: OtlpSectionConfig) -> PluginResult<CoreOpenInferenceConfig> {
1361 let transport = match section.transport.as_str() {
1362 "http_binary" => OpenInferenceTransport::HttpBinary,
1363 "grpc" => OpenInferenceTransport::Grpc,
1364 other => {
1365 return Err(PluginError::InvalidConfig(format!(
1366 "OpenInference transport must be 'http_binary' or 'grpc', got {other:?}"
1367 )));
1368 }
1369 };
1370 let mut config = CoreOpenInferenceConfig::new()
1371 .with_transport(transport)
1372 .with_service_name(section.service_name)
1373 .with_timeout(Duration::from_millis(section.timeout_millis));
1374
1375 if let Some(endpoint) = section.endpoint {
1376 config = config.with_endpoint(endpoint);
1377 }
1378 if let Some(namespace) = section.service_namespace {
1379 config = config.with_service_namespace(namespace);
1380 }
1381 if let Some(version) = section.service_version {
1382 config = config.with_service_version(version);
1383 }
1384 if let Some(scope) = section.instrumentation_scope {
1385 config = config.with_instrumentation_scope(scope);
1386 }
1387 for (key, value) in section.headers {
1388 config = config.with_header(key, value);
1389 }
1390 for (key, value) in section.resource_attributes {
1391 config = config.with_resource_attribute(key, value);
1392 }
1393 Ok(config)
1394}
1395
1396fn parse_observability_config(
1397 plugin_config: &Map<String, Json>,
1398) -> PluginResult<ObservabilityConfig> {
1399 serde_json::from_value(Json::Object(plugin_config.clone())).map_err(|err| {
1400 PluginError::InvalidConfig(format!("invalid observability plugin config: {err}"))
1401 })
1402}
1403
1404fn validate_observability_plugin_config(
1405 plugin_config: &Map<String, Json>,
1406) -> Vec<ConfigDiagnostic> {
1407 let config = match parse_observability_config(plugin_config) {
1408 Ok(config) => config,
1409 Err(err) => {
1410 return vec![ConfigDiagnostic {
1411 level: DiagnosticLevel::Error,
1412 code: "observability.invalid_plugin_config".to_string(),
1413 component: Some(OBSERVABILITY_PLUGIN_KIND.to_string()),
1414 field: None,
1415 message: err.to_string(),
1416 }];
1417 }
1418 };
1419
1420 let mut diagnostics = vec![];
1421 validate_top_level_observability_fields(&mut diagnostics, &config.policy, plugin_config);
1422 validate_version(&mut diagnostics, &config.policy, config.version);
1423 validate_policy_fields(&mut diagnostics, &config.policy, plugin_config);
1424 validate_observability_section_fields(&mut diagnostics, &config.policy, plugin_config);
1425 validate_observability_section_values(&mut diagnostics, &config);
1426
1427 diagnostics
1428}
1429
1430fn validate_top_level_observability_fields(
1431 diagnostics: &mut Vec<ConfigDiagnostic>,
1432 policy: &ConfigPolicy,
1433 plugin_config: &Map<String, Json>,
1434) {
1435 validate_unknown_fields(
1436 diagnostics,
1437 policy,
1438 Some(OBSERVABILITY_PLUGIN_KIND.to_string()),
1439 plugin_config,
1440 &[
1441 "version",
1442 "atof",
1443 "atif",
1444 "opentelemetry",
1445 "openinference",
1446 "policy",
1447 ],
1448 );
1449}
1450
1451fn validate_observability_section_fields(
1452 diagnostics: &mut Vec<ConfigDiagnostic>,
1453 policy: &ConfigPolicy,
1454 plugin_config: &Map<String, Json>,
1455) {
1456 validate_section_fields(
1457 diagnostics,
1458 policy,
1459 plugin_config,
1460 "atof",
1461 &[
1462 "enabled",
1463 "output_directory",
1464 "filename",
1465 "mode",
1466 "endpoints",
1467 ],
1468 );
1469 validate_section_fields(
1470 diagnostics,
1471 policy,
1472 plugin_config,
1473 "atif",
1474 &[
1475 "enabled",
1476 "agent_name",
1477 "agent_version",
1478 "model_name",
1479 "tool_definitions",
1480 "extra",
1481 "output_directory",
1482 "filename_template",
1483 "storage",
1484 ],
1485 );
1486 validate_section_fields(
1487 diagnostics,
1488 policy,
1489 plugin_config,
1490 "opentelemetry",
1491 &[
1492 "enabled",
1493 "transport",
1494 "endpoint",
1495 "headers",
1496 "resource_attributes",
1497 "service_name",
1498 "service_namespace",
1499 "service_version",
1500 "instrumentation_scope",
1501 "timeout_millis",
1502 ],
1503 );
1504 validate_section_fields(
1505 diagnostics,
1506 policy,
1507 plugin_config,
1508 "openinference",
1509 &[
1510 "enabled",
1511 "transport",
1512 "endpoint",
1513 "headers",
1514 "resource_attributes",
1515 "service_name",
1516 "service_namespace",
1517 "service_version",
1518 "instrumentation_scope",
1519 "timeout_millis",
1520 ],
1521 );
1522}
1523
1524fn validate_observability_section_values(
1525 diagnostics: &mut Vec<ConfigDiagnostic>,
1526 config: &ObservabilityConfig,
1527) {
1528 if let Some(section) = &config.atof {
1529 validate_atof_section(diagnostics, &config.policy, section);
1530 }
1531 if let Some(section) = &config.atif {
1532 validate_atif_section(diagnostics, &config.policy, section);
1533 }
1534 if let Some(section) = &config.opentelemetry {
1535 validate_opentelemetry_section(diagnostics, &config.policy, section);
1536 }
1537 if let Some(section) = &config.openinference {
1538 validate_openinference_section(diagnostics, &config.policy, section);
1539 }
1540}
1541
1542fn validate_atof_section(
1543 diagnostics: &mut Vec<ConfigDiagnostic>,
1544 policy: &ConfigPolicy,
1545 section: &AtofSectionConfig,
1546) {
1547 validate_atof_values(diagnostics, policy, section);
1548 validate_atof_feature_support(diagnostics, policy, section);
1549}
1550
1551#[cfg(target_arch = "wasm32")]
1552fn validate_atof_feature_support(
1553 diagnostics: &mut Vec<ConfigDiagnostic>,
1554 policy: &ConfigPolicy,
1555 section: &AtofSectionConfig,
1556) {
1557 if section.enabled {
1558 push_policy_diag(
1559 diagnostics,
1560 policy.unsupported_value,
1561 "observability.unsupported_value",
1562 Some("atof".to_string()),
1563 Some("enabled".to_string()),
1564 "ATOF file export is not supported on WebAssembly".to_string(),
1565 );
1566 }
1567 if section.enabled && !section.endpoints.is_empty() {
1568 push_policy_diag(
1569 diagnostics,
1570 policy.unsupported_value,
1571 "observability.unsupported_value",
1572 Some("atof".to_string()),
1573 Some("endpoints".to_string()),
1574 "ATOF streaming endpoints are not supported on WebAssembly".to_string(),
1575 );
1576 }
1577}
1578
1579#[cfg(all(not(feature = "atof-streaming"), not(target_arch = "wasm32")))]
1580fn validate_atof_feature_support(
1581 diagnostics: &mut Vec<ConfigDiagnostic>,
1582 policy: &ConfigPolicy,
1583 section: &AtofSectionConfig,
1584) {
1585 if section.enabled && !section.endpoints.is_empty() {
1586 push_policy_diag(
1587 diagnostics,
1588 policy.unsupported_value,
1589 "observability.unsupported_value",
1590 Some("atof".to_string()),
1591 Some("endpoints".to_string()),
1592 "ATOF streaming endpoints are not enabled in this build".to_string(),
1593 );
1594 }
1595}
1596
1597#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
1598fn validate_atof_feature_support(
1599 _diagnostics: &mut Vec<ConfigDiagnostic>,
1600 _policy: &ConfigPolicy,
1601 _section: &AtofSectionConfig,
1602) {
1603}
1604
1605fn validate_atif_section(
1606 diagnostics: &mut Vec<ConfigDiagnostic>,
1607 policy: &ConfigPolicy,
1608 section: &AtifSectionConfig,
1609) {
1610 validate_atif_values(diagnostics, policy, section);
1611 validate_atif_file_export_support(diagnostics, policy, section);
1612 validate_atif_storage_support(diagnostics, policy, section);
1613}
1614
1615#[cfg(target_arch = "wasm32")]
1616fn validate_atif_file_export_support(
1617 diagnostics: &mut Vec<ConfigDiagnostic>,
1618 policy: &ConfigPolicy,
1619 section: &AtifSectionConfig,
1620) {
1621 if section.enabled {
1622 push_policy_diag(
1623 diagnostics,
1624 policy.unsupported_value,
1625 "observability.unsupported_value",
1626 Some("atif".to_string()),
1627 Some("enabled".to_string()),
1628 "ATIF file export is not supported on WebAssembly".to_string(),
1629 );
1630 }
1631}
1632
1633#[cfg(not(target_arch = "wasm32"))]
1634fn validate_atif_file_export_support(
1635 _diagnostics: &mut Vec<ConfigDiagnostic>,
1636 _policy: &ConfigPolicy,
1637 _section: &AtifSectionConfig,
1638) {
1639}
1640
1641#[cfg(not(all(feature = "object-store", not(target_arch = "wasm32"))))]
1642fn validate_atif_storage_support(
1643 diagnostics: &mut Vec<ConfigDiagnostic>,
1644 policy: &ConfigPolicy,
1645 section: &AtifSectionConfig,
1646) {
1647 if section.enabled && !section.storage.is_empty() {
1648 push_policy_diag(
1649 diagnostics,
1650 policy.unsupported_value,
1651 "observability.feature_disabled",
1652 Some("atif".to_string()),
1653 Some("storage".to_string()),
1654 "ATIF storage support is not enabled in this build".to_string(),
1655 );
1656 }
1657}
1658
1659#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
1660fn validate_atif_storage_support(
1661 _diagnostics: &mut Vec<ConfigDiagnostic>,
1662 _policy: &ConfigPolicy,
1663 _section: &AtifSectionConfig,
1664) {
1665}
1666
1667fn validate_opentelemetry_section(
1668 diagnostics: &mut Vec<ConfigDiagnostic>,
1669 policy: &ConfigPolicy,
1670 section: &OtlpSectionConfig,
1671) {
1672 validate_otlp_values(diagnostics, policy, "opentelemetry", section);
1673 validate_opentelemetry_feature_support(diagnostics, policy, section);
1674}
1675
1676#[cfg(not(feature = "otel"))]
1677fn validate_opentelemetry_feature_support(
1678 diagnostics: &mut Vec<ConfigDiagnostic>,
1679 policy: &ConfigPolicy,
1680 section: &OtlpSectionConfig,
1681) {
1682 if section.enabled {
1683 push_policy_diag(
1684 diagnostics,
1685 policy.unsupported_value,
1686 "observability.feature_disabled",
1687 Some("opentelemetry".to_string()),
1688 Some("enabled".to_string()),
1689 "OpenTelemetry support is not enabled in this build".to_string(),
1690 );
1691 }
1692}
1693
1694#[cfg(feature = "otel")]
1695fn validate_opentelemetry_feature_support(
1696 _diagnostics: &mut Vec<ConfigDiagnostic>,
1697 _policy: &ConfigPolicy,
1698 _section: &OtlpSectionConfig,
1699) {
1700}
1701
1702fn validate_openinference_section(
1703 diagnostics: &mut Vec<ConfigDiagnostic>,
1704 policy: &ConfigPolicy,
1705 section: &OtlpSectionConfig,
1706) {
1707 validate_otlp_values(diagnostics, policy, "openinference", section);
1708 validate_openinference_feature_support(diagnostics, policy, section);
1709}
1710
1711#[cfg(not(feature = "openinference"))]
1712fn validate_openinference_feature_support(
1713 diagnostics: &mut Vec<ConfigDiagnostic>,
1714 policy: &ConfigPolicy,
1715 section: &OtlpSectionConfig,
1716) {
1717 if section.enabled {
1718 push_policy_diag(
1719 diagnostics,
1720 policy.unsupported_value,
1721 "observability.feature_disabled",
1722 Some("openinference".to_string()),
1723 Some("enabled".to_string()),
1724 "OpenInference support is not enabled in this build".to_string(),
1725 );
1726 }
1727}
1728
1729#[cfg(feature = "openinference")]
1730fn validate_openinference_feature_support(
1731 _diagnostics: &mut Vec<ConfigDiagnostic>,
1732 _policy: &ConfigPolicy,
1733 _section: &OtlpSectionConfig,
1734) {
1735}
1736
1737fn validate_version(diagnostics: &mut Vec<ConfigDiagnostic>, policy: &ConfigPolicy, version: u32) {
1738 if version != 1 {
1739 push_policy_diag(
1740 diagnostics,
1741 policy.unsupported_value,
1742 "observability.unsupported_config_version",
1743 Some(OBSERVABILITY_PLUGIN_KIND.to_string()),
1744 Some("version".to_string()),
1745 format!("observability config version {version} is unsupported"),
1746 );
1747 }
1748}
1749
1750fn validate_policy_fields(
1751 diagnostics: &mut Vec<ConfigDiagnostic>,
1752 policy: &ConfigPolicy,
1753 plugin_config: &Map<String, Json>,
1754) {
1755 if let Some(policy_json) = plugin_config.get("policy").and_then(Json::as_object) {
1756 validate_unknown_fields(
1757 diagnostics,
1758 policy,
1759 Some("policy".to_string()),
1760 policy_json,
1761 &["unknown_component", "unknown_field", "unsupported_value"],
1762 );
1763 }
1764}
1765
1766fn validate_section_fields(
1767 diagnostics: &mut Vec<ConfigDiagnostic>,
1768 policy: &ConfigPolicy,
1769 plugin_config: &Map<String, Json>,
1770 section: &str,
1771 known_fields: &[&str],
1772) {
1773 if let Some(section_json) = plugin_config.get(section).and_then(Json::as_object) {
1774 validate_unknown_fields(
1775 diagnostics,
1776 policy,
1777 Some(section.to_string()),
1778 section_json,
1779 known_fields,
1780 );
1781 }
1782}
1783
1784fn validate_atof_values(
1785 diagnostics: &mut Vec<ConfigDiagnostic>,
1786 policy: &ConfigPolicy,
1787 section: &AtofSectionConfig,
1788) {
1789 if AtofExporterMode::parse(§ion.mode).is_none() {
1790 push_policy_diag(
1791 diagnostics,
1792 policy.unsupported_value,
1793 "observability.unsupported_value",
1794 Some("atof".to_string()),
1795 Some("mode".to_string()),
1796 "ATOF mode must be 'append' or 'overwrite'".to_string(),
1797 );
1798 }
1799 for (index, endpoint) in section.endpoints.iter().enumerate() {
1800 validate_atof_endpoint_values(diagnostics, policy, index, endpoint);
1801 }
1802}
1803
1804fn validate_atof_endpoint_values(
1805 diagnostics: &mut Vec<ConfigDiagnostic>,
1806 policy: &ConfigPolicy,
1807 index: usize,
1808 endpoint: &AtofEndpointSectionConfig,
1809) {
1810 if endpoint.url.trim().is_empty() {
1811 push_policy_diag(
1812 diagnostics,
1813 policy.unsupported_value,
1814 "observability.unsupported_value",
1815 Some("atof".to_string()),
1816 Some(format!("endpoints[{index}].url")),
1817 format!("ATOF endpoints[{index}].url must be non-empty"),
1818 );
1819 } else if !is_valid_atof_endpoint_url(&endpoint.url) {
1820 push_policy_diag(
1821 diagnostics,
1822 policy.unsupported_value,
1823 "observability.unsupported_value",
1824 Some("atof".to_string()),
1825 Some(format!("endpoints[{index}].url")),
1826 format!("ATOF endpoints[{index}].url must be a valid URL"),
1827 );
1828 }
1829 if AtofEndpointTransport::parse(&endpoint.transport).is_none() {
1830 push_policy_diag(
1831 diagnostics,
1832 policy.unsupported_value,
1833 "observability.unsupported_value",
1834 Some("atof".to_string()),
1835 Some(format!("endpoints[{index}].transport")),
1836 format!(
1837 "ATOF endpoints[{index}].transport must be 'http_post', 'websocket', or 'ndjson'"
1838 ),
1839 );
1840 }
1841 if endpoint.timeout_millis == 0 {
1842 push_policy_diag(
1843 diagnostics,
1844 policy.unsupported_value,
1845 "observability.unsupported_value",
1846 Some("atof".to_string()),
1847 Some(format!("endpoints[{index}].timeout_millis")),
1848 format!("ATOF endpoints[{index}].timeout_millis must be greater than 0"),
1849 );
1850 }
1851}
1852
1853#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
1854fn is_valid_atof_endpoint_url(url: &str) -> bool {
1855 reqwest::Url::parse(url).is_ok()
1856}
1857
1858#[cfg(any(not(feature = "atof-streaming"), target_arch = "wasm32"))]
1859fn is_valid_atof_endpoint_url(_url: &str) -> bool {
1860 true
1861}
1862
1863fn validate_atif_values(
1864 diagnostics: &mut Vec<ConfigDiagnostic>,
1865 policy: &ConfigPolicy,
1866 section: &AtifSectionConfig,
1867) {
1868 if !section.filename_template.contains("{session_id}") {
1869 push_policy_diag(
1870 diagnostics,
1871 policy.unsupported_value,
1872 "observability.unsupported_value",
1873 Some("atif".to_string()),
1874 Some("filename_template".to_string()),
1875 "ATIF filename_template must contain '{session_id}'".to_string(),
1876 );
1877 }
1878 for (index, storage) in section.storage.iter().enumerate() {
1879 validate_atif_storage_values(diagnostics, policy, index, storage);
1880 }
1881}
1882
1883fn validate_atif_storage_values(
1884 diagnostics: &mut Vec<ConfigDiagnostic>,
1885 policy: &ConfigPolicy,
1886 index: usize,
1887 storage: &AtifStorageConfig,
1888) {
1889 match storage {
1890 AtifStorageConfig::Http(http) => {
1891 validate_atif_http_endpoint(
1892 diagnostics,
1893 policy,
1894 &format!("storage[{index}].endpoint"),
1895 &http.endpoint,
1896 );
1897 if http.timeout_millis == 0 {
1898 push_policy_diag(
1899 diagnostics,
1900 policy.unsupported_value,
1901 "observability.unsupported_value",
1902 Some("atif".to_string()),
1903 Some(format!("storage[{index}].timeout_millis")),
1904 format!("ATIF storage[{index}].timeout_millis must be positive"),
1905 );
1906 }
1907 for (header, value) in &http.headers {
1908 validate_atif_http_header(
1909 diagnostics,
1910 policy,
1911 &format!("storage[{index}].headers.{header}"),
1912 header,
1913 value,
1914 );
1915 }
1916 for (header, var_name) in &http.header_env {
1917 validate_atif_http_header_name(
1918 diagnostics,
1919 policy,
1920 &format!("storage[{index}].header_env.{header}"),
1921 header,
1922 );
1923 validate_atif_storage_env_var(
1924 diagnostics,
1925 policy,
1926 &format!("storage[{index}].header_env.{header}"),
1927 Some(var_name.as_str()),
1928 );
1929 }
1930 }
1931 AtifStorageConfig::S3(s3) => {
1932 if s3.bucket.trim().is_empty() {
1933 push_policy_diag(
1934 diagnostics,
1935 policy.unsupported_value,
1936 "observability.unsupported_value",
1937 Some("atif".to_string()),
1938 Some(format!("storage[{index}].bucket")),
1939 format!("ATIF storage[{index}].bucket must be non-empty"),
1940 );
1941 }
1942 validate_atif_storage_env_var(
1943 diagnostics,
1944 policy,
1945 &format!("storage[{index}].secret_access_key_var"),
1946 s3.secret_access_key_var.as_deref(),
1947 );
1948 validate_atif_storage_env_var(
1949 diagnostics,
1950 policy,
1951 &format!("storage[{index}].session_token_var"),
1952 s3.session_token_var.as_deref(),
1953 );
1954 }
1955 }
1956}
1957
1958fn validate_atif_http_header(
1959 diagnostics: &mut Vec<ConfigDiagnostic>,
1960 policy: &ConfigPolicy,
1961 field: &str,
1962 header: &str,
1963 _value: &str,
1964) {
1965 validate_atif_http_header_name(diagnostics, policy, field, header);
1966 #[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
1967 if let Err(err) = reqwest::header::HeaderValue::from_str(_value) {
1968 push_policy_diag(
1969 diagnostics,
1970 policy.unsupported_value,
1971 "observability.unsupported_value",
1972 Some("atif".to_string()),
1973 Some(field.to_string()),
1974 format!("ATIF {field} value is invalid: {err}"),
1975 );
1976 }
1977}
1978
1979fn validate_atif_http_header_name(
1980 diagnostics: &mut Vec<ConfigDiagnostic>,
1981 policy: &ConfigPolicy,
1982 field: &str,
1983 header: &str,
1984) {
1985 #[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
1986 let is_valid = reqwest::header::HeaderName::from_bytes(header.as_bytes()).is_ok();
1987 #[cfg(not(all(feature = "object-store", not(target_arch = "wasm32"))))]
1988 let is_valid = !header.trim().is_empty() && header.trim() == header;
1989 if !is_valid {
1990 push_policy_diag(
1991 diagnostics,
1992 policy.unsupported_value,
1993 "observability.unsupported_value",
1994 Some("atif".to_string()),
1995 Some(field.to_string()),
1996 format!("ATIF {field} header name '{header}' is invalid"),
1997 );
1998 }
1999}
2000
2001fn validate_atif_http_endpoint(
2002 diagnostics: &mut Vec<ConfigDiagnostic>,
2003 policy: &ConfigPolicy,
2004 field: &str,
2005 endpoint: &str,
2006) {
2007 let trimmed = endpoint.trim();
2008 let mut is_valid = !trimmed.is_empty() && trimmed == endpoint;
2009 #[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2010 {
2011 is_valid = is_valid
2012 && reqwest::Url::parse(endpoint)
2013 .map(|url| matches!(url.scheme(), "http" | "https") && url.host_str().is_some())
2014 .unwrap_or(false);
2015 }
2016 #[cfg(not(all(feature = "object-store", not(target_arch = "wasm32"))))]
2017 {
2018 let valid_scheme = trimmed.starts_with("http://") || trimmed.starts_with("https://");
2019 let has_host = trimmed
2020 .split_once("://")
2021 .map(|(_, rest)| !rest.is_empty() && !rest.starts_with('/'))
2022 .unwrap_or(false);
2023 is_valid = is_valid && valid_scheme && has_host;
2024 }
2025 if !is_valid {
2026 push_policy_diag(
2027 diagnostics,
2028 policy.unsupported_value,
2029 "observability.unsupported_value",
2030 Some("atif".to_string()),
2031 Some(field.to_string()),
2032 format!("ATIF {field} must be a valid http:// or https:// URL"),
2033 );
2034 }
2035}
2036
2037fn validate_atif_storage_env_var(
2038 diagnostics: &mut Vec<ConfigDiagnostic>,
2039 policy: &ConfigPolicy,
2040 field: &str,
2041 var_name: Option<&str>,
2042) {
2043 let Some(var_name) = var_name else {
2044 return;
2045 };
2046 let trimmed = var_name.trim();
2047 if trimmed.is_empty() {
2048 push_policy_diag(
2049 diagnostics,
2050 policy.unsupported_value,
2051 "observability.unsupported_value",
2052 Some("atif".to_string()),
2053 Some(field.to_string()),
2054 format!("ATIF {field} must be the name of an environment variable, not empty"),
2055 );
2056 return;
2057 }
2058 if trimmed != var_name {
2059 push_policy_diag(
2060 diagnostics,
2061 policy.unsupported_value,
2062 "observability.unsupported_value",
2063 Some("atif".to_string()),
2064 Some(field.to_string()),
2065 format!("ATIF {field} must not have surrounding whitespace; got '{var_name}'"),
2066 );
2067 return;
2068 }
2069 match std::env::var(var_name) {
2070 Ok(value) if !value.is_empty() => {}
2071 Ok(_) => {
2072 push_policy_diag(
2073 diagnostics,
2074 policy.unsupported_value,
2075 "observability.unsupported_value",
2076 Some("atif".to_string()),
2077 Some(field.to_string()),
2078 format!(
2079 "ATIF {field}='{var_name}' references an environment variable that is set but empty"
2080 ),
2081 );
2082 }
2083 Err(_) => {
2084 push_policy_diag(
2085 diagnostics,
2086 policy.unsupported_value,
2087 "observability.unsupported_value",
2088 Some("atif".to_string()),
2089 Some(field.to_string()),
2090 format!(
2091 "ATIF {field}='{var_name}' references an environment variable that is not set"
2092 ),
2093 );
2094 }
2095 }
2096}
2097
2098fn validate_otlp_values(
2099 diagnostics: &mut Vec<ConfigDiagnostic>,
2100 policy: &ConfigPolicy,
2101 section_name: &str,
2102 section: &OtlpSectionConfig,
2103) {
2104 if !matches!(section.transport.as_str(), "http_binary" | "grpc") {
2105 push_policy_diag(
2106 diagnostics,
2107 policy.unsupported_value,
2108 "observability.unsupported_value",
2109 Some(section_name.to_string()),
2110 Some("transport".to_string()),
2111 format!("{section_name} transport must be 'http_binary' or 'grpc'"),
2112 );
2113 }
2114}
2115
2116fn validate_unknown_fields(
2117 diagnostics: &mut Vec<ConfigDiagnostic>,
2118 policy: &ConfigPolicy,
2119 component: Option<String>,
2120 config: &Map<String, Json>,
2121 known_fields: &[&str],
2122) {
2123 for field in config.keys() {
2124 if !known_fields.contains(&field.as_str()) {
2125 push_policy_diag(
2126 diagnostics,
2127 policy.unknown_field,
2128 "observability.unknown_field",
2129 component.clone(),
2130 Some(field.clone()),
2131 format!(
2132 "field '{}' is not recognized for '{}'",
2133 field,
2134 component.as_deref().unwrap_or("unknown")
2135 ),
2136 );
2137 }
2138 }
2139}
2140
2141fn push_policy_diag(
2142 diagnostics: &mut Vec<ConfigDiagnostic>,
2143 behavior: UnsupportedBehavior,
2144 code: &str,
2145 component: Option<String>,
2146 field: Option<String>,
2147 message: String,
2148) {
2149 let level = match behavior {
2150 UnsupportedBehavior::Ignore => return,
2151 UnsupportedBehavior::Warn => DiagnosticLevel::Warning,
2152 UnsupportedBehavior::Error => DiagnosticLevel::Error,
2153 };
2154 diagnostics.push(ConfigDiagnostic {
2155 level,
2156 code: code.to_string(),
2157 component,
2158 field,
2159 message,
2160 });
2161}
2162
2163fn observability_registration_error(error: impl std::fmt::Display) -> PluginError {
2164 PluginError::RegistrationFailed(error.to_string())
2165}
2166
2167fn default_observability_config_version() -> u32 {
2168 1
2169}
2170
2171fn default_atof_mode() -> String {
2172 "append".to_string()
2173}
2174
2175fn default_atof_endpoint_transport() -> String {
2176 "http_post".to_string()
2177}
2178
2179fn default_agent_name() -> String {
2180 "NeMo Relay".to_string()
2181}
2182
2183fn default_agent_version() -> String {
2184 env!("CARGO_PKG_VERSION").to_string()
2185}
2186
2187fn default_model_name() -> String {
2188 "unknown".to_string()
2189}
2190
2191fn default_atif_filename_template() -> String {
2192 "nemo-relay-atif-{session_id}.json".to_string()
2193}
2194
2195fn default_otlp_transport() -> String {
2196 "http_binary".to_string()
2197}
2198
2199fn default_service_name() -> String {
2200 "nemo-relay".to_string()
2201}
2202
2203fn default_timeout_millis() -> u64 {
2204 3_000
2205}
2206
2207fn default_output_directory() -> PathBuf {
2208 std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
2209}
2210
2211#[cfg(not(all(feature = "object-store", not(target_arch = "wasm32"))))]
2212struct AtifRemoteStorage;
2213
2214#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2222struct AtifRemoteStorage {
2223 sender: std::sync::mpsc::Sender<AtifUploadRequest>,
2224 key_prefix: String,
2225}
2226
2227#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2228struct AtifUploadRequest {
2229 key: String,
2230 filename: String,
2231 session_id: String,
2232 payload: Vec<u8>,
2233 reply: std::sync::mpsc::Sender<std::io::Result<()>>,
2234}
2235
2236#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2237#[derive(Clone)]
2238struct HttpUploadConfig {
2239 endpoint: String,
2240 headers: HashMap<String, String>,
2241 timeout: Duration,
2242}
2243
2244#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2245#[derive(Default)]
2246struct S3BuilderOverrides {
2247 access_key_id: Option<String>,
2248 secret_access_key: Option<String>,
2249 session_token: Option<String>,
2250 region: Option<String>,
2251 endpoint_url: Option<String>,
2252 allow_http: Option<bool>,
2253}
2254
2255#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2256impl S3BuilderOverrides {
2257 fn resolve(index: usize, s3: &S3StorageConfig) -> std::io::Result<Self> {
2258 Ok(Self {
2259 access_key_id: s3.access_key_id.clone(),
2260 secret_access_key: resolve_env_var_field(
2261 &format!("storage[{index}].secret_access_key_var"),
2262 s3.secret_access_key_var.as_deref(),
2263 )?,
2264 session_token: resolve_env_var_field(
2265 &format!("storage[{index}].session_token_var"),
2266 s3.session_token_var.as_deref(),
2267 )?,
2268 region: s3.region.clone(),
2269 endpoint_url: s3.endpoint_url.clone(),
2270 allow_http: s3.allow_http,
2271 })
2272 }
2273
2274 fn apply(
2275 self,
2276 mut builder: object_store::aws::AmazonS3Builder,
2277 ) -> object_store::aws::AmazonS3Builder {
2278 if let Some(value) = self.access_key_id {
2279 builder = builder.with_access_key_id(value);
2280 }
2281 if let Some(value) = self.secret_access_key {
2282 builder = builder.with_secret_access_key(value);
2283 }
2284 if let Some(value) = self.session_token {
2285 builder = builder.with_token(value);
2286 }
2287 if let Some(value) = self.region {
2288 builder = builder.with_region(value);
2289 }
2290 if let Some(value) = self.endpoint_url {
2291 builder = builder.with_endpoint(value);
2292 }
2293 if let Some(value) = self.allow_http {
2294 builder = builder.with_allow_http(value);
2295 }
2296 builder
2297 }
2298}
2299
2300#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2301fn resolve_env_var_field(field: &str, var_name: Option<&str>) -> std::io::Result<Option<String>> {
2302 let Some(var_name) = var_name else {
2303 return Ok(None);
2304 };
2305 if var_name.trim().is_empty() || var_name.trim() != var_name {
2306 return Err(std::io::Error::other(format!(
2307 "ATIF {field} must be the name of an environment variable, not '{var_name}'"
2308 )));
2309 }
2310 match std::env::var(var_name) {
2311 Ok(value) if !value.is_empty() => Ok(Some(value)),
2312 Ok(_) => Err(std::io::Error::other(format!(
2313 "ATIF {field}='{var_name}' references an environment variable that is set but empty"
2314 ))),
2315 Err(_) => Err(std::io::Error::other(format!(
2316 "ATIF {field}='{var_name}' references an environment variable that is not set"
2317 ))),
2318 }
2319}
2320
2321#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2322impl AtifRemoteStorage {
2323 fn from_config(index: usize, config: &AtifStorageConfig) -> std::io::Result<Self> {
2324 match config {
2325 AtifStorageConfig::Http(http) => Self::build_http(index, http),
2326 AtifStorageConfig::S3(s3) => Self::build_s3(index, s3),
2327 }
2328 }
2329
2330 fn build_http(index: usize, http: &HttpStorageConfig) -> std::io::Result<Self> {
2331 let upload_config = HttpUploadConfig::resolve(index, http)?;
2332 let (req_tx, req_rx) = std::sync::mpsc::channel::<AtifUploadRequest>();
2333 let (ready_tx, ready_rx) = std::sync::mpsc::channel::<std::io::Result<()>>();
2334
2335 std::thread::Builder::new()
2336 .name("nemo-relay-atif-storage".to_string())
2337 .spawn(move || {
2338 install_rustls_crypto_provider();
2339 let runtime = match tokio::runtime::Builder::new_current_thread()
2340 .enable_all()
2341 .build()
2342 {
2343 Ok(rt) => rt,
2344 Err(err) => {
2345 let _ = ready_tx.send(Err(std::io::Error::other(format!(
2346 "failed to build ATIF storage runtime: {err}"
2347 ))));
2348 return;
2349 }
2350 };
2351 let client = match reqwest::Client::builder()
2352 .timeout(upload_config.timeout)
2353 .build()
2354 {
2355 Ok(client) => client,
2356 Err(err) => {
2357 let _ = ready_tx.send(Err(std::io::Error::other(format!(
2358 "failed to build HTTP client for ATIF storage[{}]: {err}",
2359 index
2360 ))));
2361 return;
2362 }
2363 };
2364 if ready_tx.send(Ok(())).is_err() {
2365 return;
2366 }
2367 drop(ready_tx);
2368
2369 while let Ok(request) = req_rx.recv() {
2370 let result = runtime.block_on(post_atif_http(
2371 &client,
2372 &upload_config,
2373 request.filename,
2374 request.session_id,
2375 request.payload,
2376 ));
2377 let _ = request.reply.send(result);
2378 }
2379 })
2380 .map_err(|err| {
2381 std::io::Error::other(format!("failed to spawn ATIF storage thread: {err}"))
2382 })?;
2383
2384 match ready_rx.recv() {
2385 Ok(Ok(())) => Ok(Self {
2386 sender: req_tx,
2387 key_prefix: String::new(),
2388 }),
2389 Ok(Err(err)) => Err(err),
2390 Err(_) => Err(std::io::Error::other(
2391 "ATIF storage thread exited before signalling readiness",
2392 )),
2393 }
2394 }
2395
2396 fn build_s3(index: usize, s3: &S3StorageConfig) -> std::io::Result<Self> {
2397 let bucket = s3.bucket.clone();
2398 let key_prefix = normalize_storage_key_prefix(s3.key_prefix.as_deref());
2399 let overrides = S3BuilderOverrides::resolve(index, s3)?;
2400
2401 let (req_tx, req_rx) = std::sync::mpsc::channel::<AtifUploadRequest>();
2402 let (ready_tx, ready_rx) = std::sync::mpsc::channel::<std::io::Result<()>>();
2403
2404 std::thread::Builder::new()
2405 .name("nemo-relay-atif-storage".to_string())
2406 .spawn(move || {
2407 install_rustls_crypto_provider();
2408 let runtime = match tokio::runtime::Builder::new_current_thread()
2409 .enable_all()
2410 .build()
2411 {
2412 Ok(rt) => rt,
2413 Err(err) => {
2414 let _ = ready_tx.send(Err(std::io::Error::other(format!(
2415 "failed to build ATIF storage runtime: {err}"
2416 ))));
2417 return;
2418 }
2419 };
2420 let store = match overrides
2421 .apply(object_store::aws::AmazonS3Builder::from_env())
2422 .with_bucket_name(&bucket)
2423 .build()
2424 {
2425 Ok(store) => Arc::new(store) as Arc<dyn object_store::ObjectStore>,
2426 Err(err) => {
2427 let _ = ready_tx.send(Err(std::io::Error::other(format!(
2428 "failed to build S3 client for bucket '{bucket}': {err}"
2429 ))));
2430 return;
2431 }
2432 };
2433 if ready_tx.send(Ok(())).is_err() {
2434 return;
2435 }
2436 drop(ready_tx);
2437
2438 while let Ok(request) = req_rx.recv() {
2439 let result = runtime.block_on(async {
2440 use object_store::ObjectStoreExt as _;
2441 store
2442 .put(
2443 &object_store::path::Path::from(request.key.clone()),
2444 object_store::PutPayload::from(request.payload),
2445 )
2446 .await
2447 .map(|_| ())
2448 .map_err(|err| {
2449 std::io::Error::other(format!(
2450 "S3 upload to '{}' failed: {err}",
2451 request.key
2452 ))
2453 })
2454 });
2455 let _ = request.reply.send(result);
2456 }
2457 })
2458 .map_err(|err| {
2459 std::io::Error::other(format!("failed to spawn ATIF storage thread: {err}"))
2460 })?;
2461
2462 match ready_rx.recv() {
2463 Ok(Ok(())) => Ok(Self {
2464 sender: req_tx,
2465 key_prefix,
2466 }),
2467 Ok(Err(err)) => Err(err),
2468 Err(_) => Err(std::io::Error::other(
2469 "ATIF storage thread exited before signalling readiness",
2470 )),
2471 }
2472 }
2473
2474 fn put(&self, filename: &str, session_id: &str, payload: &[u8]) -> std::io::Result<()> {
2475 let key = format!("{}{}", self.key_prefix, filename);
2476 let (reply_tx, reply_rx) = std::sync::mpsc::channel();
2477 self.sender
2478 .send(AtifUploadRequest {
2479 key,
2480 filename: filename.to_string(),
2481 session_id: session_id.to_string(),
2482 payload: payload.to_vec(),
2483 reply: reply_tx,
2484 })
2485 .map_err(|_| std::io::Error::other("ATIF storage thread is not running"))?;
2486 reply_rx
2487 .recv()
2488 .map_err(|_| std::io::Error::other("ATIF storage thread dropped the upload reply"))?
2489 }
2490}
2491
2492#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2493impl HttpUploadConfig {
2494 fn resolve(index: usize, http: &HttpStorageConfig) -> std::io::Result<Self> {
2495 let endpoint = http.endpoint.trim();
2496 if endpoint.is_empty() || endpoint != http.endpoint {
2497 return Err(std::io::Error::other(format!(
2498 "ATIF storage[{index}].endpoint must be non-empty and must not have surrounding whitespace"
2499 )));
2500 }
2501 let parsed = reqwest::Url::parse(endpoint).map_err(|err| {
2502 std::io::Error::other(format!(
2503 "ATIF storage[{index}].endpoint must be a valid URL: {err}"
2504 ))
2505 })?;
2506 if !matches!(parsed.scheme(), "http" | "https") || parsed.host_str().is_none() {
2507 return Err(std::io::Error::other(format!(
2508 "ATIF storage[{index}].endpoint must be a valid http:// or https:// URL"
2509 )));
2510 }
2511 if http.timeout_millis == 0 {
2512 return Err(std::io::Error::other(format!(
2513 "ATIF storage[{index}].timeout_millis must be positive"
2514 )));
2515 }
2516
2517 let mut headers = http.headers.clone();
2518 for (header, var_name) in &http.header_env {
2519 let value = resolve_env_var_field(
2520 &format!("storage[{index}].header_env.{header}"),
2521 Some(var_name.as_str()),
2522 )?
2523 .expect("resolve_env_var_field returns Some when var_name is Some");
2524 headers.insert(header.clone(), value);
2525 }
2526 validate_http_headers(index, &headers)?;
2527
2528 Ok(Self {
2529 endpoint: parsed.to_string(),
2530 headers,
2531 timeout: Duration::from_millis(http.timeout_millis),
2532 })
2533 }
2534}
2535
2536#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2537fn validate_http_headers(index: usize, headers: &HashMap<String, String>) -> std::io::Result<()> {
2538 for (header, value) in headers {
2539 reqwest::header::HeaderName::from_bytes(header.as_bytes()).map_err(|err| {
2540 std::io::Error::other(format!(
2541 "ATIF storage[{index}] header name '{header}' is invalid: {err}"
2542 ))
2543 })?;
2544 reqwest::header::HeaderValue::from_str(value).map_err(|err| {
2545 std::io::Error::other(format!(
2546 "ATIF storage[{index}] value for header '{header}' is invalid: {err}"
2547 ))
2548 })?;
2549 }
2550 Ok(())
2551}
2552
2553#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2554async fn post_atif_http(
2555 client: &reqwest::Client,
2556 config: &HttpUploadConfig,
2557 filename: String,
2558 session_id: String,
2559 payload: Vec<u8>,
2560) -> std::io::Result<()> {
2561 let mut request = client.post(&config.endpoint);
2562 for (header, value) in &config.headers {
2563 request = request.header(header.as_str(), value.as_str());
2564 }
2565 let response = request
2566 .header(reqwest::header::CONTENT_TYPE, "application/json")
2567 .header("x-nemo-relay-atif-filename", filename.clone())
2568 .header("x-nemo-relay-atif-session-id", session_id)
2569 .body(payload)
2570 .send()
2571 .await
2572 .map_err(|err| {
2573 std::io::Error::other(format!(
2574 "HTTP ATIF upload to '{}' failed: {err}",
2575 config.endpoint
2576 ))
2577 })?;
2578 if response.status().is_success() {
2579 Ok(())
2580 } else {
2581 Err(std::io::Error::other(format!(
2582 "HTTP ATIF upload to '{}' for '{}' failed with status {}",
2583 config.endpoint,
2584 filename,
2585 response.status()
2586 )))
2587 }
2588}
2589
2590#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2591fn install_rustls_crypto_provider() {
2592 let _ = rustls::crypto::ring::default_provider().install_default();
2593}
2594
2595#[cfg(all(feature = "object-store", not(target_arch = "wasm32")))]
2596fn normalize_storage_key_prefix(raw: Option<&str>) -> String {
2597 let trimmed = raw.unwrap_or("").trim();
2598 if trimmed.is_empty() {
2599 return String::new();
2600 }
2601 if trimmed.ends_with('/') {
2602 trimmed.to_string()
2603 } else {
2604 format!("{trimmed}/")
2605 }
2606}
2607
2608#[cfg(test)]
2609#[path = "../../tests/unit/observability/plugin_component_tests.rs"]
2610mod tests;