Skip to content

feat: record provider request and response sizes #712

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bk-nvidia wants to merge 2 commits into
base: main
Choose a base branch
from
+202 −3
Open

feat: record provider request and response sizes #712

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9c6ba67
feat: record provider request and response sizes
bk-nvidia Aug 5, 2026
94c56c2
fix: emit response size before stream errors
bk-nvidia Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 36 additions & 2 deletions crates/cli/src/gateway/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
}

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -593,20 +596,24 @@ 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;
}
}
}
}
Err(error) => {
emit_provider_body_size("gateway.provider.response", body_size_bytes);
yield Err(FlowError::Internal(error.to_string()));
return;
}
Expand All @@ -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)
}
Expand Down Expand Up @@ -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)
Expand All @@ -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,
Expand Down Expand Up @@ -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)
}
Expand Down
161 changes: 160 additions & 1 deletion crates/cli/tests/coverage/shared/gateway_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::<Value>));
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!(
Expand Down Expand Up @@ -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::<Value>));
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();
}
}),
Comment on lines +1041 to +1045

@coderabbitai coderabbitai Bot Aug 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Isolate global subscriber capture in both tests.

Each callback accepts every gateway.provider.response event and overwrites the captured value. Concurrent gateway tests can therefore make either assertion inspect another request’s body size.

  • crates/cli/tests/coverage/shared/gateway_tests.rs#L1041-L1045: filter on a test-specific event identifier, or serialize access to the global subscriber.
  • crates/cli/tests/coverage/shared/gateway_tests.rs#L1089-L1093: apply the same isolation mechanism.

As per path instructions, tests must cover “cross-request isolation where relevant.”

📍 Affects 1 file
  • crates/cli/tests/coverage/shared/gateway_tests.rs#L1041-L1045 (this comment)
  • crates/cli/tests/coverage/shared/gateway_tests.rs#L1089-L1093
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/cli/tests/coverage/shared/gateway_tests.rs` around lines 1041 - 1045,
Isolate global subscriber capture in both gateway tests by updating the
callbacks at crates/cli/tests/coverage/shared/gateway_tests.rs:1041-1045 and
:1089-1093 to accept only each test’s unique event identifier, or serialize
subscriber access. Preserve the existing response-event capture and assertions
while ensuring concurrent requests cannot overwrite another test’s captured body
size; include cross-request isolation coverage.

Source: Path instructions

)
.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::<Value>));
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]
Expand Down Expand Up @@ -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::<String, Value>::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();
Comment thread
coderabbitai[bot] marked this conversation as resolved.
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 {
Expand Down Expand Up @@ -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();
Expand All @@ -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]
Expand Down
6 changes: 6 additions & 0 deletions docs/nemo-relay-cli/basic-usage.mdx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading