diff --git a/crates/core/src/observability/atof.rs b/crates/core/src/observability/atof.rs index 45f50613d..f6f312c36 100644 --- a/crates/core/src/observability/atof.rs +++ b/crates/core/src/observability/atof.rs @@ -644,6 +644,14 @@ struct AtofEndpointWorker { transport: AtofEndpointTransport, } +/// Validated endpoint settings with headers resolved at activation. +#[cfg(feature = "atof-streaming")] +#[derive(Debug)] +struct ActivatedAtofEndpoint { + config: AtofEndpointConfig, + headers: reqwest::header::HeaderMap, +} + impl AtofEndpointWorker { fn enqueue(&self, raw_json: String) { let _ = self.sender.send(EndpointMessage::Event(raw_json)); @@ -711,13 +719,13 @@ fn start_endpoint_workers(configs: &[AtofEndpointConfig]) -> Result Result { - validate_endpoint_config(&config)?; - let timeout = Duration::from_millis(config.timeout_millis); - let transport = config.transport; + let endpoint = validate_endpoint_config(config)?; + let timeout = Duration::from_millis(endpoint.config.timeout_millis); + let transport = endpoint.config.transport; let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); std::thread::Builder::new() .name(format!("nemo-relay-atof-endpoint-{index}")) - .spawn(move || run_endpoint_worker(index, config, rx)) + .spawn(move || run_endpoint_worker(index, endpoint, rx)) .map_err(|error| AtofExporterError::InvalidEndpoint(error.to_string()))?; log::info!( target: "nemo_relay.observability", @@ -745,7 +753,7 @@ fn start_endpoint_worker(index: usize, config: AtofEndpointConfig) -> Result Result<()> { +fn validate_endpoint_config(config: AtofEndpointConfig) -> Result { if config.url.trim().is_empty() { return Err(AtofExporterError::InvalidEndpoint( "endpoint url must be non-empty".to_string(), @@ -771,8 +779,8 @@ fn validate_endpoint_config(config: &AtofEndpointConfig) -> Result<()> { url.scheme() ))); } - resolved_header_map(&config.headers, &config.header_env)?; - Ok(()) + let headers = resolved_header_map(&config.headers, &config.header_env)?; + Ok(ActivatedAtofEndpoint { config, headers }) } #[cfg(feature = "atof-streaming")] @@ -816,7 +824,7 @@ fn resolved_header_map( #[cfg(feature = "atof-streaming")] fn run_endpoint_worker( index: usize, - config: AtofEndpointConfig, + endpoint: ActivatedAtofEndpoint, rx: tokio::sync::mpsc::UnboundedReceiver, ) { let runtime = match tokio::runtime::Builder::new_current_thread() @@ -837,10 +845,10 @@ fn run_endpoint_worker( } }; runtime.block_on(async move { - match config.transport { - AtofEndpointTransport::HttpPost => run_http_post_endpoint(index, config, rx).await, - AtofEndpointTransport::Websocket => run_websocket_endpoint(index, config, rx).await, - AtofEndpointTransport::Ndjson => run_ndjson_endpoint(index, config, rx).await, + match endpoint.config.transport { + AtofEndpointTransport::HttpPost => run_http_post_endpoint(index, endpoint, rx).await, + AtofEndpointTransport::Websocket => run_websocket_endpoint(index, endpoint, rx).await, + AtofEndpointTransport::Ndjson => run_ndjson_endpoint(index, endpoint, rx).await, } }); } @@ -848,29 +856,12 @@ fn run_endpoint_worker( #[cfg(feature = "atof-streaming")] async fn run_http_post_endpoint( index: usize, - config: AtofEndpointConfig, + endpoint: ActivatedAtofEndpoint, mut rx: tokio::sync::mpsc::UnboundedReceiver, ) { let client = match reqwest::Client::builder() - .timeout(Duration::from_millis(config.timeout_millis)) - .default_headers( - match resolved_header_map(&config.headers, &config.header_env) { - Ok(headers) => headers, - Err(_) => { - log::error!( - target: "nemo_relay.observability", - event = "endpoint_disabled", - exporter = "atof", - endpoint_index = index, - transport = "http_post", - reason = "invalid_headers"; - "ATOF endpoint disabled" - ); - drain_closed(rx).await; - return; - } - }, - ) + .timeout(Duration::from_millis(endpoint.config.timeout_millis)) + .default_headers(endpoint.headers) .build() { Ok(client) => client, @@ -892,9 +883,9 @@ async fn run_http_post_endpoint( while let Some(message) = rx.recv().await { match message { EndpointMessage::Event(raw_json) => { - let body = format!("{}\n", endpoint_event_json(&config, raw_json)); + let body = format!("{}\n", endpoint_event_json(&endpoint.config, raw_json)); let result = client - .post(&config.url) + .post(&endpoint.config.url) .header(reqwest::header::CONTENT_TYPE, "application/x-ndjson") .body(body) .send() @@ -941,12 +932,12 @@ async fn run_http_post_endpoint( #[cfg(feature = "atof-streaming")] async fn run_websocket_endpoint( index: usize, - config: AtofEndpointConfig, + endpoint: ActivatedAtofEndpoint, mut rx: tokio::sync::mpsc::UnboundedReceiver, ) { let mut pending = std::collections::VecDeque::new(); let mut retry = WebSocketRetryState::default(); - let mut socket = match connect_websocket(&config).await { + let mut socket = match connect_websocket(&endpoint).await { Ok(socket) => { retry.record_recovered(index); Some(socket) @@ -959,21 +950,36 @@ async fn run_websocket_endpoint( while let Some(message) = rx.recv().await { match message { EndpointMessage::Event(raw_json) => { - pending.push_back(endpoint_event_json(&config, raw_json)); - let _ = - drain_websocket_pending(index, &config, &mut socket, &mut pending, &mut retry) - .await; + pending.push_back(endpoint_event_json(&endpoint.config, raw_json)); + let _ = drain_websocket_pending( + index, + &endpoint, + &mut socket, + &mut pending, + &mut retry, + ) + .await; } EndpointMessage::Flush(done) => { - let _ = - drain_websocket_pending(index, &config, &mut socket, &mut pending, &mut retry) - .await; + let _ = drain_websocket_pending( + index, + &endpoint, + &mut socket, + &mut pending, + &mut retry, + ) + .await; let _ = done.send(()); } EndpointMessage::Close(done) => { - let _ = - drain_websocket_pending(index, &config, &mut socket, &mut pending, &mut retry) - .await; + let _ = drain_websocket_pending( + index, + &endpoint, + &mut socket, + &mut pending, + &mut retry, + ) + .await; if let Some(mut ws) = socket.take() { let _ = ws.close(None).await; } @@ -1046,16 +1052,16 @@ impl WebSocketRetryState { #[cfg(feature = "atof-streaming")] async fn drain_websocket_pending( index: usize, - config: &AtofEndpointConfig, + endpoint: &ActivatedAtofEndpoint, socket: &mut Option, pending: &mut std::collections::VecDeque, retry: &mut WebSocketRetryState, ) -> bool { - let timeout = Duration::from_millis(config.timeout_millis); + let timeout = Duration::from_millis(endpoint.config.timeout_millis); let attempts_before = retry.attempts; match tokio::time::timeout( timeout, - drain_websocket_pending_inner(index, config, socket, pending, retry), + drain_websocket_pending_inner(index, endpoint, socket, pending, retry), ) .await { @@ -1072,14 +1078,14 @@ async fn drain_websocket_pending( #[cfg(feature = "atof-streaming")] async fn drain_websocket_pending_inner( index: usize, - config: &AtofEndpointConfig, + endpoint: &ActivatedAtofEndpoint, socket: &mut Option, pending: &mut std::collections::VecDeque, retry: &mut WebSocketRetryState, ) -> bool { while let Some(raw_json) = pending.front().cloned() { if socket.is_none() { - match connect_websocket(config).await { + match connect_websocket(endpoint).await { Ok(ws) => { *socket = Some(ws); retry.record_recovered(index); @@ -1116,23 +1122,22 @@ async fn drain_websocket_pending_inner( #[cfg(feature = "atof-streaming")] async fn connect_websocket( - config: &AtofEndpointConfig, + endpoint: &ActivatedAtofEndpoint, ) -> std::result::Result { - let mut request = config + let mut request = endpoint + .config .url .as_str() .into_client_request() .map_err(|error| error.to_string())?; - let headers = resolved_header_map(&config.headers, &config.header_env) - .map_err(|error| error.to_string())?; - for (name, value) in headers { + for (name, value) in endpoint.headers.clone() { let Some(name) = name else { continue; }; request.headers_mut().insert(name, value); } tokio::time::timeout( - Duration::from_millis(config.timeout_millis), + Duration::from_millis(endpoint.config.timeout_millis), tokio_tungstenite::connect_async(request), ) .await @@ -1144,10 +1149,10 @@ async fn connect_websocket( #[cfg(feature = "atof-streaming")] async fn run_ndjson_endpoint( index: usize, - config: AtofEndpointConfig, + endpoint: ActivatedAtofEndpoint, mut rx: tokio::sync::mpsc::UnboundedReceiver, ) { - let client = match build_ndjson_client(&config) { + let client = match build_ndjson_client(&endpoint) { Ok(client) => client, Err(_) => { log::error!( @@ -1165,7 +1170,7 @@ async fn run_ndjson_endpoint( }; let (body_tx, body) = ndjson_body_channel(); - let url = config.url.clone(); + let url = endpoint.config.url.clone(); let request = tokio::spawn(async move { client .post(url) @@ -1174,13 +1179,15 @@ async fn run_ndjson_endpoint( .send() .await }); - let close_timeout = Duration::from_millis(config.timeout_millis); + let close_timeout = Duration::from_millis(endpoint.config.timeout_millis); while let Some(message) = rx.recv().await { match message { - EndpointMessage::Event(raw_json) => { - send_ndjson_event(index, &body_tx, endpoint_event_json(&config, raw_json)) - } + EndpointMessage::Event(raw_json) => send_ndjson_event( + index, + &body_tx, + endpoint_event_json(&endpoint.config, raw_json), + ), EndpointMessage::Flush(done) => send_ndjson_flush(index, &body_tx, done), EndpointMessage::Close(done) => { drop(body_tx); @@ -1193,13 +1200,11 @@ async fn run_ndjson_endpoint( #[cfg(feature = "atof-streaming")] fn build_ndjson_client( - config: &AtofEndpointConfig, + endpoint: &ActivatedAtofEndpoint, ) -> std::result::Result { - let headers = resolved_header_map(&config.headers, &config.header_env) - .map_err(|error| format!("disabled: {error}"))?; reqwest::Client::builder() - .connect_timeout(Duration::from_millis(config.timeout_millis)) - .default_headers(headers) + .connect_timeout(Duration::from_millis(endpoint.config.timeout_millis)) + .default_headers(endpoint.headers.clone()) .build() .map_err(|error| format!("client build failed: {error}")) } diff --git a/crates/core/tests/unit/observability/atof_tests.rs b/crates/core/tests/unit/observability/atof_tests.rs index 7bf2b8c14..d18be0989 100644 --- a/crates/core/tests/unit/observability/atof_tests.rs +++ b/crates/core/tests/unit/observability/atof_tests.rs @@ -1378,7 +1378,7 @@ fn invalid_endpoint_scheme_errors_cleanly() { fn endpoint_validation_rejects_empty_timeout_and_invalid_headers() { let mut headers = std::collections::HashMap::new(); headers.insert("x-test".to_string(), "ok".to_string()); - validate_endpoint_config(&AtofEndpointConfig { + validate_endpoint_config(AtofEndpointConfig { url: "http://127.0.0.1:9/events".into(), transport: AtofEndpointTransport::HttpPost, headers: headers.clone(), @@ -1422,7 +1422,7 @@ fn endpoint_validation_rejects_empty_timeout_and_invalid_headers() { field_name_policy: AtofEndpointFieldNamePolicy::Preserve, }; assert!( - validate_endpoint_config(&empty_url) + validate_endpoint_config(empty_url) .unwrap_err() .to_string() .contains("endpoint url must be non-empty") @@ -1437,7 +1437,7 @@ fn endpoint_validation_rejects_empty_timeout_and_invalid_headers() { field_name_policy: AtofEndpointFieldNamePolicy::Preserve, }; assert!( - validate_endpoint_config(&zero_timeout) + validate_endpoint_config(zero_timeout) .unwrap_err() .to_string() .contains("timeout_millis") @@ -1451,7 +1451,7 @@ fn endpoint_validation_rejects_empty_timeout_and_invalid_headers() { bad_header_value.insert("x-test".to_string(), "bad\nvalue".to_string()); assert!(resolved_header_map(&bad_header_value, &Default::default()).is_err()); assert!( - build_ndjson_client(&AtofEndpointConfig { + validate_endpoint_config(AtofEndpointConfig { url: "http://127.0.0.1:9/events".into(), transport: AtofEndpointTransport::Ndjson, headers: bad_header_value, @@ -1459,8 +1459,33 @@ fn endpoint_validation_rejects_empty_timeout_and_invalid_headers() { timeout_millis: 1, field_name_policy: AtofEndpointFieldNamePolicy::Preserve, }) - .unwrap_err() - .contains("disabled") + .is_err() + ); +} + +#[test] +#[cfg(feature = "atof-streaming")] +fn endpoint_activation_snapshots_header_env() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + let variable = format!("NEMO_RELAY_TEST_ATOF_HEADER_ENV_{}", std::process::id()); + // SAFETY: The test mutex serializes environment access for this process. + unsafe { std::env::set_var(&variable, "Bearer relay-499") }; + + let endpoint = validate_endpoint_config(AtofEndpointConfig { + url: "http://127.0.0.1:9/events".into(), + transport: AtofEndpointTransport::HttpPost, + headers: std::collections::HashMap::new(), + header_env: std::collections::HashMap::from([("authorization".into(), variable.clone())]), + timeout_millis: 1, + field_name_policy: AtofEndpointFieldNamePolicy::Preserve, + }) + .unwrap(); + + // SAFETY: The activated endpoint must not read this environment variable again. + unsafe { std::env::remove_var(&variable) }; + assert_eq!( + endpoint.headers.get("authorization").unwrap(), + "Bearer relay-499" ); } @@ -1536,8 +1561,11 @@ fn http_endpoint_worker_acknowledges_flush_close_and_logs_http_errors() { tokio::runtime::Runtime::new().unwrap().block_on(async { run_http_post_endpoint( 0, - AtofEndpointConfig::new(url, AtofEndpointTransport::HttpPost) - .with_timeout_millis(5_000), + validate_endpoint_config( + AtofEndpointConfig::new(url, AtofEndpointTransport::HttpPost) + .with_timeout_millis(5_000), + ) + .unwrap(), rx, ) .await; @@ -1560,46 +1588,6 @@ fn http_endpoint_worker_acknowledges_flush_close_and_logs_http_errors() { server.join().unwrap(); } -#[test] -#[cfg(feature = "atof-streaming")] -fn http_endpoint_worker_disables_invalid_headers_and_drains_control_messages() { - enable_operational_logs(); - let mut headers = std::collections::HashMap::new(); - headers.insert("bad header".to_string(), "ok".to_string()); - let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); - let worker = std::thread::spawn(move || { - tokio::runtime::Runtime::new().unwrap().block_on(async { - run_http_post_endpoint( - 0, - AtofEndpointConfig { - url: "http://127.0.0.1:9/events".into(), - transport: AtofEndpointTransport::HttpPost, - headers, - header_env: std::collections::HashMap::new(), - timeout_millis: 1, - field_name_policy: AtofEndpointFieldNamePolicy::Preserve, - }, - rx, - ) - .await; - }); - }); - - tx.send(EndpointMessage::Event("{\"kind\":\"mark\"}".into())) - .unwrap(); - let (flush_tx, flush_rx) = std::sync::mpsc::channel(); - tx.send(EndpointMessage::Flush(flush_tx)).unwrap(); - let (close_tx, close_rx) = std::sync::mpsc::channel(); - tx.send(EndpointMessage::Close(close_tx)).unwrap(); - flush_rx - .recv_timeout(std::time::Duration::from_secs(1)) - .unwrap(); - close_rx - .recv_timeout(std::time::Duration::from_secs(1)) - .unwrap(); - worker.join().unwrap(); -} - #[test] #[cfg(feature = "atof-streaming")] fn http_endpoint_worker_reports_request_transport_failure() { @@ -1609,11 +1597,14 @@ fn http_endpoint_worker_reports_request_transport_failure() { tokio::runtime::Runtime::new().unwrap().block_on(async { run_http_post_endpoint( 0, - AtofEndpointConfig::new( - "http://127.0.0.1:0/events", - AtofEndpointTransport::HttpPost, + validate_endpoint_config( + AtofEndpointConfig::new( + "http://127.0.0.1:0/events", + AtofEndpointTransport::HttpPost, + ) + .with_timeout_millis(100), ) - .with_timeout_millis(100), + .unwrap(), rx, ) .await; @@ -1630,65 +1621,28 @@ fn http_endpoint_worker_reports_request_transport_failure() { worker.join().unwrap(); } -#[test] -#[cfg(feature = "atof-streaming")] -fn ndjson_endpoint_worker_disables_invalid_headers_and_drains_control_messages() { - enable_operational_logs(); - let mut headers = std::collections::HashMap::new(); - headers.insert("bad header".to_string(), "ok".to_string()); - let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); - let worker = std::thread::spawn(move || { - tokio::runtime::Runtime::new().unwrap().block_on(async { - run_ndjson_endpoint( - 0, - AtofEndpointConfig { - url: "http://127.0.0.1:9/events".into(), - transport: AtofEndpointTransport::Ndjson, - headers, - header_env: std::collections::HashMap::new(), - timeout_millis: 1, - field_name_policy: AtofEndpointFieldNamePolicy::Preserve, - }, - rx, - ) - .await; - }); - }); - - let (flush_tx, flush_rx) = std::sync::mpsc::channel(); - tx.send(EndpointMessage::Flush(flush_tx)).unwrap(); - let (close_tx, close_rx) = std::sync::mpsc::channel(); - tx.send(EndpointMessage::Close(close_tx)).unwrap(); - flush_rx - .recv_timeout(std::time::Duration::from_secs(1)) - .unwrap(); - close_rx - .recv_timeout(std::time::Duration::from_secs(1)) - .unwrap(); - worker.join().unwrap(); -} - #[test] #[cfg(feature = "atof-streaming")] fn websocket_helpers_cover_invalid_headers_and_timeout_reconnect_path() { enable_operational_logs(); - let mut headers = std::collections::HashMap::new(); - headers.insert("bad header".to_string(), "ok".to_string()); - let config = AtofEndpointConfig { + let endpoint = validate_endpoint_config(AtofEndpointConfig { url: "ws://127.0.0.1:9/events".into(), transport: AtofEndpointTransport::Websocket, - headers, + headers: std::collections::HashMap::new(), header_env: std::collections::HashMap::new(), timeout_millis: 1, field_name_policy: AtofEndpointFieldNamePolicy::Preserve, - }; + }) + .unwrap(); tokio::runtime::Runtime::new().unwrap().block_on(async { - assert!(connect_websocket(&config).await.is_err()); + assert!(connect_websocket(&endpoint).await.is_err()); let mut socket = None; let mut pending = std::collections::VecDeque::from(["{\"kind\":\"mark\"}".to_string()]); let mut retry = WebSocketRetryState::default(); - assert!(!drain_websocket_pending(0, &config, &mut socket, &mut pending, &mut retry).await); + assert!( + !drain_websocket_pending(0, &endpoint, &mut socket, &mut pending, &mut retry).await + ); assert_eq!(pending.len(), 1); assert_eq!(retry.attempts, 1); }); diff --git a/crates/core/tests/unit/observability/plugin_component_tests.rs b/crates/core/tests/unit/observability/plugin_component_tests.rs index 6c42e717d..b086fe3ea 100644 --- a/crates/core/tests/unit/observability/plugin_component_tests.rs +++ b/crates/core/tests/unit/observability/plugin_component_tests.rs @@ -37,7 +37,14 @@ fn temp_dir(prefix: &str) -> PathBuf { } #[cfg(feature = "atof-streaming")] -fn start_http_capture_server(expected_requests: usize) -> (String, Arc>>) { +#[derive(Clone, Debug)] +struct HttpCapture { + headers: String, + body: String, +} + +#[cfg(feature = "atof-streaming")] +fn start_http_capture_server(expected_requests: usize) -> (String, Arc>>) { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let url = format!("http://{}", listener.local_addr().unwrap()); let captures = Arc::new(Mutex::new(Vec::new())); @@ -65,10 +72,10 @@ fn start_http_capture_server(expected_requests: usize) -> (String, Arc (String, std::thread::JoinH } #[cfg(feature = "atof-streaming")] -fn wait_for_captures(captures: &Arc>>, expected: usize) -> Vec { +fn wait_for_captures(captures: &Arc>>, expected: usize) -> Vec { for _ in 0..100 { let snapshot = captures.lock().unwrap().clone(); if snapshot.len() >= expected { @@ -1115,13 +1122,64 @@ fn atof_stream_sinks_fan_out_and_teardown_all_workers() { for captures in [&first_captures, &second_captures] { let bodies = wait_for_captures(captures, 3); assert_eq!(bodies.len(), 3, "captured bodies: {bodies:?}"); - let events = bodies.join(""); + let events = bodies + .iter() + .map(|capture| capture.body.as_str()) + .collect::(); assert!(events.contains("\"scope_category\":\"start\"")); assert!(events.contains("\"name\":\"checkpoint\"")); assert!(events.contains("\"scope_category\":\"end\"")); } } +#[test] +#[cfg(feature = "atof-streaming")] +fn atof_stream_sink_header_env_is_snapshotted_at_activation() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + reset_runtime(); + let variable = format!("NEMO_RELAY_TEST_ATOF_HEADER_ENV_{}", std::process::id()); + // SAFETY: The test mutex serializes environment access for this process. + unsafe { std::env::set_var(&variable, "Bearer relay-499") }; + let (url, captures) = start_http_capture_server(3); + + let config = plugin_config(json!({ + "atof": { + "enabled": true, + "sinks": [{ + "type": "stream", + "url": url, + "transport": "http_post", + "header_env": {"authorization": variable.clone()} + }] + } + })); + futures::executor::block_on(initialize_plugins_exact(config)).unwrap(); + // SAFETY: The active endpoint must use the header value captured at activation. + unsafe { std::env::remove_var(&variable) }; + + let agent = push_agent("atof-header-env-agent"); + crate::api::scope::event( + crate::api::scope::EmitMarkEventParams::builder() + .name("checkpoint") + .parent(&agent) + .data(json!({"step": 1})) + .build(), + ) + .unwrap(); + pop(&agent); + clear_plugin_configuration().unwrap(); + + let captures = wait_for_captures(&captures, 3); + assert_eq!(captures.len(), 3, "captured requests: {captures:?}"); + for capture in captures { + assert!(capture.headers.lines().any(|line| { + line.split_once(':').is_some_and(|(name, value)| { + name.eq_ignore_ascii_case("authorization") && value.trim() == "Bearer relay-499" + }) + })); + } +} + #[test] #[cfg(all(feature = "atof-streaming", feature = "object-store"))] fn atif_remote_storage_validates_s3_configuration_and_http_access_outcomes() { diff --git a/python/tests/test_observability_plugin.py b/python/tests/test_observability_plugin.py index 70ff55e7a..734f05acf 100644 --- a/python/tests/test_observability_plugin.py +++ b/python/tests/test_observability_plugin.py @@ -5,7 +5,10 @@ from __future__ import annotations +import http.server import json +import threading +import time import typing import pytest @@ -29,6 +32,55 @@ from pathlib import Path +class _AtofCaptureServer(http.server.ThreadingHTTPServer): + requests: list[tuple[dict[str, str], bytes]] + request_event: threading.Event + + +class _AtofCaptureHandler(http.server.BaseHTTPRequestHandler): + def do_POST(self) -> None: # noqa: N802 + content_length = int(self.headers.get("content-length", "0")) + server = typing.cast(_AtofCaptureServer, self.server) + server.requests.append((dict(self.headers.items()), self.rfile.read(content_length))) + server.request_event.set() + self.send_response(200) + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: # noqa: ARG002 + return + + +class _AtofCapture: + server: "_AtofCaptureServer" + thread: threading.Thread + + def __enter__(self) -> _AtofCapture: + self.server = _AtofCaptureServer(("127.0.0.1", 0), _AtofCaptureHandler) + self.server.requests = [] + self.server.request_event = threading.Event() + self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) + self.thread.start() + return self + + def __exit__(self, *args: object) -> None: + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=1) + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.server.server_port}" + + def wait_for_requests(self, expected: int, timeout: float = 5.0) -> list[tuple[dict[str, str], bytes]]: + deadline = time.monotonic() + timeout + while len(self.server.requests) < expected: + remaining = deadline - time.monotonic() + assert remaining > 0, f"timed out waiting for {expected} ATOF requests" + self.server.request_event.wait(remaining) + self.server.request_event.clear() + return self.server.requests + + class TestObservabilityConfigHelpers: def test_defaults_and_component_wrapper(self): assert AtofConfig().to_dict() == {"enabled": False} @@ -131,6 +183,48 @@ def test_atof_endpoint_alias_preserves_positional_transport(self): assert endpoint.transport == "websocket" assert endpoint.name is None + async def test_atof_stream_sink_snapshots_header_env( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ): + variable = "NEMO_RELAY_TEST_ATOF_HEADER_ENV" + credential = "Bearer relay-499" + monkeypatch.setenv(variable, credential) + + with _AtofCapture() as capture: + config = ObservabilityConfig( + atof=AtofConfig( + enabled=True, + sinks=[ + AtofStreamSinkConfig( + url=capture.url, + transport="http_post", + header_env={"authorization": variable}, + ) + ], + ) + ) + report = await plugin.initialize(plugin.PluginConfig(components=[ComponentSpec(config)])) + assert report["diagnostics"] == [] + monkeypatch.delenv(variable) + + try: + with scope.scope("python-header-env-agent", ScopeType.Agent) as handle: + scope.event("python-header-env-mark", handle=handle, data={"step": 1}) + finally: + plugin.clear() + requests = capture.wait_for_requests(3) + + assert len(requests) == 3 + payload = b"".join(body for _, body in requests).decode() + assert '"scope_category":"start"' in payload + assert '"name":"python-header-env-mark"' in payload + assert '"scope_category":"end"' in payload + for headers, _ in requests: + authorization = next(value for name, value in headers.items() if name.lower() == "authorization") + assert authorization == credential + assert credential not in json.dumps(report) + assert credential not in caplog.text + def test_http_storage_config_serializes_headers(self): s3 = S3StorageConfig(bucket="archive") http = HttpStorageConfig(