diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index 76feccd4c..de7401d20 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -182,6 +182,7 @@ async fn run_unmanaged_gateway( let status = response.status(); let headers = response_headers(response.headers()); let bytes = response.bytes().await?; + emit_provider_body_size("gateway.provider.response", bytes.len()); build_response(status, headers, Body::from(bytes)) } @@ -368,6 +369,7 @@ fn build_buffered_func( ); } }; + emit_provider_body_size("gateway.provider.response", bytes.len()); if !status.is_success() { if retry_aware { return Err(FlowError::Upstream(http_failure( @@ -563,6 +565,7 @@ fn build_streaming_func( ); } }; + emit_provider_body_size("gateway.provider.response", bytes.len()); if retry_aware { return Err(FlowError::Upstream(http_failure( status, @@ -593,13 +596,16 @@ fn sse_json_stream(response: reqwest::Response) -> LlmJsonStream { let mut decoder = SseEventDecoder::new(); let mut bytes = response.bytes_stream(); let stream = stream! { + let mut body_size_bytes = 0usize; while let Some(chunk) = bytes.next().await { match chunk { Ok(buffer) => { + body_size_bytes = body_size_bytes.saturating_add(buffer.len()); for result in decoder.push_bytes_results(&buffer) { match result { Ok(event) => yield Ok(event.data), Err(error) => { + emit_provider_body_size("gateway.provider.response", body_size_bytes); yield Err(error); return; } @@ -607,6 +613,7 @@ fn sse_json_stream(response: reqwest::Response) -> LlmJsonStream { } } Err(error) => { + emit_provider_body_size("gateway.provider.response", body_size_bytes); yield Err(FlowError::Internal(error.to_string())); return; } @@ -615,8 +622,13 @@ fn sse_json_stream(response: reqwest::Response) -> LlmJsonStream { match decoder.finish() { Ok(Some(event)) => yield Ok(event.data), Ok(None) => {} - Err(error) => yield Err(error), + Err(error) => { + emit_provider_body_size("gateway.provider.response", body_size_bytes); + yield Err(error); + return; + } } + emit_provider_body_size("gateway.provider.response", body_size_bytes); }; LlmJsonStream::new(stream) } @@ -804,6 +816,7 @@ async fn forward_upstream_request( url, forwarding.source_route, ); + emit_provider_body_size("gateway.provider.request", effective.body_bytes.len()); let configured_auth_header = forwarding.configured_auth_header(effective.target_route); let mut upstream = http .request(method.clone(), &effective.url) @@ -826,6 +839,15 @@ async fn forward_upstream_request( upstream.send().await } +fn emit_provider_body_size(name: &'static str, body_size_bytes: usize) { + let _ = nemo_relay::api::scope::event( + nemo_relay::api::scope::EmitMarkEventParams::builder() + .name(name) + .data(serde_json::json!({"body_size_bytes": body_size_bytes})) + .build(), + ); +} + #[derive(Clone)] struct EffectiveUpstreamRequest { body_bytes: Bytes, @@ -1113,9 +1135,21 @@ async fn passthrough_streaming( let headers = response_headers(response.headers()); let mut bytes = response.bytes_stream(); let body = Body::from_stream(stream! { + let mut body_size_bytes = 0usize; while let Some(chunk) = bytes.next().await { - yield chunk; + match chunk { + Ok(buffer) => { + body_size_bytes = body_size_bytes.saturating_add(buffer.len()); + yield Ok(buffer); + } + Err(error) => { + emit_provider_body_size("gateway.provider.response", body_size_bytes); + yield Err(error); + return; + } + } } + emit_provider_body_size("gateway.provider.response", body_size_bytes); }); build_response(status, headers, body) } diff --git a/crates/cli/tests/coverage/shared/gateway_tests.rs b/crates/cli/tests/coverage/shared/gateway_tests.rs index a775ddf93..dde3d0121 100644 --- a/crates/cli/tests/coverage/shared/gateway_tests.rs +++ b/crates/cli/tests/coverage/shared/gateway_tests.rs @@ -974,6 +974,19 @@ fn structured_upstream_failure_classification_matches_retry_policy() { #[tokio::test] async fn sse_json_stream_yields_valid_event_before_later_batch_error() { + let subscriber_name = "gateway-streaming-provider-response-body-size-test"; + let _ = nemo_relay::api::subscriber::deregister_subscriber(subscriber_name); + let captured_data = Arc::new(Mutex::new(None::)); + let captured = captured_data.clone(); + nemo_relay::api::subscriber::register_subscriber( + subscriber_name, + Arc::new(move |event| { + if event.name() == "gateway.provider.response" { + *captured.lock().unwrap() = event.data().cloned(); + } + }), + ) + .unwrap(); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); let sse_body = concat!( @@ -1009,6 +1022,119 @@ async fn sse_json_stream_yields_valid_event_before_later_batch_error() { assert!(stream.next().await.is_none()); server.await.unwrap(); + nemo_relay::api::subscriber::flush_subscribers().unwrap(); + nemo_relay::api::subscriber::deregister_subscriber(subscriber_name).unwrap(); + assert_eq!( + captured_data.lock().unwrap().as_ref().unwrap()["body_size_bytes"], + json!(sse_body.len()) + ); +} + +#[tokio::test] +async fn incomplete_sse_error_emits_provider_response_body_size() { + let subscriber_name = "gateway-incomplete-sse-provider-response-body-size-test"; + let _ = nemo_relay::api::subscriber::deregister_subscriber(subscriber_name); + let captured_data = Arc::new(Mutex::new(None::)); + let captured = captured_data.clone(); + nemo_relay::api::subscriber::register_subscriber( + subscriber_name, + Arc::new(move |event| { + if event.name() == "gateway.provider.response" { + *captured.lock().unwrap() = event.data().cloned(); + } + }), + ) + .unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let sse_body = "data: {not valid json"; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + sse_body.len(), + sse_body + ); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0_u8; 1024]; + let _ = socket.read(&mut request).await.unwrap(); + socket.write_all(response.as_bytes()).await.unwrap(); + }); + + let response = test_http_client() + .get(format!("http://{address}")) + .send() + .await + .unwrap(); + let mut stream = sse_json_stream(response); + + assert!(stream.next().await.unwrap().is_err()); + drop(stream); + server.await.unwrap(); + nemo_relay::api::subscriber::flush_subscribers().unwrap(); + nemo_relay::api::subscriber::deregister_subscriber(subscriber_name).unwrap(); + assert_eq!( + captured_data.lock().unwrap().as_ref().unwrap()["body_size_bytes"], + json!(sse_body.len()) + ); +} + +#[tokio::test] +async fn passthrough_transport_error_emits_provider_response_body_size() { + let subscriber_name = "gateway-passthrough-provider-response-body-size-test"; + let _ = nemo_relay::api::subscriber::deregister_subscriber(subscriber_name); + let captured_data = Arc::new(Mutex::new(None::)); + let captured = captured_data.clone(); + nemo_relay::api::subscriber::register_subscriber( + subscriber_name, + Arc::new(move |event| { + if event.name() == "gateway.provider.response" { + *captured.lock().unwrap() = event.data().cloned(); + } + }), + ) + .unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let partial_body = "partial"; + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0_u8; 1024]; + let _ = socket.read(&mut request).await.unwrap(); + socket + .write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/octet-stream\r\ncontent-length: 64\r\nconnection: close\r\n\r\n{partial_body}" + ) + .as_bytes(), + ) + .await + .unwrap(); + }); + let state = AppState::new(GatewayConfig::default()); + let prepared = PreparedGatewayRequest { + method: Method::POST, + headers: HeaderMap::new(), + path: "/v1/models".into(), + provider: ProviderRoute::OpenAiModels, + upstream_url: format!("http://{address}/v1/models"), + body_bytes: Bytes::new(), + request_json: Value::Null, + streaming: true, + authorization: crate::provider_auth::ProviderRequestAuthorization { + source_credential: crate::provider_auth::SourceCredentialDisposition::Absent, + allow_environment_provider_auth: false, + }, + }; + + let response = passthrough_streaming(state, prepared).await.unwrap(); + assert!(response.into_body().collect().await.is_err()); + server.await.unwrap(); + nemo_relay::api::subscriber::flush_subscribers().unwrap(); + nemo_relay::api::subscriber::deregister_subscriber(subscriber_name).unwrap(); + assert_eq!( + captured_data.lock().unwrap().as_ref().unwrap()["body_size_bytes"], + json!(partial_body.len()) + ); } #[tokio::test] @@ -1143,6 +1269,26 @@ async fn buffered_body_read_failure_stays_structured() { #[tokio::test] async fn buffered_invalid_json_becomes_safe_upstream_failure() { + let subscriber_name = "gateway-provider-request-body-size-test"; + let _ = nemo_relay::api::subscriber::deregister_subscriber(subscriber_name); + let captured_data = Arc::new(Mutex::new(HashMap::::new())); + let captured = captured_data.clone(); + nemo_relay::api::subscriber::register_subscriber( + subscriber_name, + Arc::new(move |event| { + if matches!( + event.name(), + "gateway.provider.request" | "gateway.provider.response" + ) && let Some(data) = event.data() + { + captured + .lock() + .unwrap() + .insert(event.name().to_string(), data.clone()); + } + }), + ) + .unwrap(); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); let server = tokio::spawn(async move { @@ -1179,9 +1325,11 @@ async fn buffered_invalid_json_becomes_safe_upstream_failure() { ); let mut request_headers = Map::new(); request_headers.insert(INTERNAL_RETRY_AWARE_HEADER.to_string(), json!("true")); + let request_content = json!({"model": "rewritten"}); + let expected_body_size = serde_json::to_vec(&request_content).unwrap().len(); let error = func(LlmRequest { headers: request_headers, - content: json!({}), + content: request_content, }) .await .unwrap_err(); @@ -1200,6 +1348,17 @@ async fn buffered_invalid_json_becomes_safe_upstream_failure() { Some(&"provider-request".to_string()) ); server.await.unwrap(); + nemo_relay::api::subscriber::flush_subscribers().unwrap(); + nemo_relay::api::subscriber::deregister_subscriber(subscriber_name).unwrap(); + let captured_data = captured_data.lock().unwrap(); + assert_eq!( + captured_data["gateway.provider.request"]["body_size_bytes"], + json!(expected_body_size) + ); + assert_eq!( + captured_data["gateway.provider.response"]["body_size_bytes"], + json!(22) + ); } #[test] diff --git a/docs/nemo-relay-cli/basic-usage.mdx b/docs/nemo-relay-cli/basic-usage.mdx index 13417c964..f1c56dd2e 100644 --- a/docs/nemo-relay-cli/basic-usage.mdx +++ b/docs/nemo-relay-cli/basic-usage.mdx @@ -49,6 +49,12 @@ payload schemas. It removes only hop-by-hop transport headers, forwards streaming responses as streams, and emits NeMo Relay LLM start and end events under the active session scope. +The gateway emits `gateway.provider.request` and `gateway.provider.response` +telemetry events containing the exact `body_size_bytes` sent to and received +from each provider. Request bytes are measured after interceptors and JSON +re-encoding; streaming response bytes are counted before SSE decoding. The +`full` and `openinference` OpenTelemetry projections export them as span events. + ## Transparent Run Use the agent shortcuts for no-install local observability. The wrapper starts