diff --git a/crates/libsy-llm-client/src/attempt.rs b/crates/libsy-llm-client/src/attempt.rs new file mode 100644 index 0000000000..55fde1fd3b --- /dev/null +++ b/crates/libsy-llm-client/src/attempt.rs @@ -0,0 +1,541 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Accounting retained until an HTTP attempt's response is consumed or abandoned. + +use std::pin::Pin; +use std::task::{Context, Poll}; + +use futures::Stream; +use http::HeaderMap; +use switchyard_protocol::{ + LlmClientError, LlmResponseChunk, LlmResponseStream, LlmResponseStreamEvent, ModelId, Usage, + WireFormat, +}; +use tracing::Span; + +use crate::observability::{llm_client_error_type, record_gen_ai_usage}; + +pub(crate) struct AttemptGuard { + span: Span, + usage: Usage, + accounting: bool, + final_usage: bool, + model_finished: bool, + finished: bool, +} + +impl AttemptGuard { + pub(crate) fn new( + model: &ModelId, + format: WireFormat, + ordinal: u64, + max_attempts: u64, + accounting: bool, + ) -> Self { + if accounting { + Span::current().record("switchyard.accounting.unit", "summary"); + Span::current().record("switchyard.accounting.attempts", ordinal); + } + let span = tracing::debug_span!(target: "libsy", "libsy.upstream_attempt", + model = %model, gen_ai.request.model = %model, wire_format = %format, + attempt = ordinal, max_attempts, retry = ordinal > 1, + openinference.span.kind = "CHAIN", + outcome = tracing::field::Empty, status_code = tracing::field::Empty, + will_retry = tracing::field::Empty, retry_delay_ms = tracing::field::Empty, + error.type = tracing::field::Empty, otel.status_code = tracing::field::Empty, + gen_ai.response.id = tracing::field::Empty, gen_ai.response.model = tracing::field::Empty, + gen_ai.usage.input_tokens = tracing::field::Empty, + gen_ai.usage.output_tokens = tracing::field::Empty, + gen_ai.usage.cache_read.input_tokens = tracing::field::Empty, + gen_ai.usage.cache_creation.input_tokens = tracing::field::Empty, + gen_ai.usage.cache_creation_1h.input_tokens = tracing::field::Empty, + gen_ai.usage.reasoning.output_tokens = tracing::field::Empty, + switchyard.usage.input_tokens = tracing::field::Empty, + switchyard.usage.input_tokens_reported = tracing::field::Empty, + switchyard.usage.total_tokens = tracing::field::Empty, + switchyard.usage.provider_cost_usd = tracing::field::Empty, + switchyard.accounting.charge_source = tracing::field::Empty, + switchyard.accounting.usage_status = tracing::field::Empty, + switchyard.accounting.unit = if accounting { "completion_attempt" } else { "diagnostic" }, + ); + Self { + span, + usage: Usage::default(), + accounting, + final_usage: false, + model_finished: false, + finished: false, + } + } + + pub(crate) fn span(&self) -> &Span { + &self.span + } + + pub(crate) fn headers(&mut self, headers: &HeaderMap) { + if let Some(cost) = headers + .get("x-litellm-response-cost") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.trim().parse::().ok()) + .filter(|value| value.is_finite() && *value >= 0.0) + { + self.usage.provider_cost = Some(cost); + self.span + .record("switchyard.accounting.charge_source", "header"); + } + } + + pub(crate) fn record_usage(&mut self, usage: &Usage, final_report: bool) { + macro_rules! update { + ($($field:ident),+ $(,)?) => {$(if usage.$field.is_some() { self.usage.$field = usage.$field; })+}; + } + update!(input_tokens, output_tokens, total_tokens, reasoning_tokens); + if let Some(cache) = &usage.cache { + let target = self.usage.cache.get_or_insert_with(Default::default); + if cache.cached_input_tokens.is_some() { + target.cached_input_tokens = cache.cached_input_tokens; + } + if cache.cache_creation_input_tokens.is_some() { + target.cache_creation_input_tokens = cache.cache_creation_input_tokens; + } + if cache.cache_creation_1h_input_tokens.is_some() { + target.cache_creation_1h_input_tokens = cache.cache_creation_1h_input_tokens; + } + } + if let Some(cost) = usage + .provider_cost + .filter(|value| value.is_finite() && *value >= 0.0) + { + self.usage.provider_cost = Some(cost); + self.span + .record("switchyard.accounting.charge_source", "provider"); + } + self.final_usage |= final_report; + } + + pub(crate) fn response(&mut self, body: &serde_json::Value, format: WireFormat) { + if let Some(id) = body.get("id").and_then(serde_json::Value::as_str) { + self.span.record("gen_ai.response.id", id); + } + if let Some(model) = body.get("model").and_then(serde_json::Value::as_str) { + self.span.record("gen_ai.response.model", model); + } + self.record_usage(&switchyard_translation::response_usage(body, format), true); + } + + pub(crate) fn observe( + &mut self, + item: &Result, + ) -> bool { + let event = match item { + Ok(event) => event, + Err(error) => { + self.error(error); + return true; + } + }; + self.observe_usage(event); + for chunk in event.normalized() { + match chunk { + LlmResponseChunk::MessageStart { id, model } => { + if let Some(id) = id { + self.span.record("gen_ai.response.id", id.as_str()); + } + if let Some(model) = model { + self.span.record("gen_ai.response.model", model.as_str()); + } + } + LlmResponseChunk::MessageStop { .. } => self.model_finished = true, + LlmResponseChunk::StreamError { .. } => { + self.finish("error", Some("provider_stream")); + return true; + } + LlmResponseChunk::DecodeError { .. } => { + self.finish("error", Some("response_translation")); + return true; + } + _ => (), + } + } + false + } + + pub(crate) fn observe_usage(&mut self, event: &LlmResponseStreamEvent) { + // One provider event can contain both its final usage and a terminal error. + let terminal_report = event.preservation().is_some_and(|raw| { + matches!( + raw.raw().get("type").and_then(serde_json::Value::as_str), + Some("response.completed" | "response.failed" | "response.incomplete") + ) + }); + for chunk in event.normalized() { + if let LlmResponseChunk::Usage(usage) = chunk { + self.record_usage(usage, terminal_report); + } + } + } + + pub(crate) fn error(&mut self, error: &LlmClientError) { + self.finish("error", Some(llm_client_error_type(error).as_ref())); + } + + pub(crate) fn success(&mut self) { + self.finish("ok", None); + } + + fn finish(&mut self, outcome: &str, error: Option<&str>) { + if self.finished { + return; + } + self.finished = true; + self.span.record("outcome", outcome); + self.span.record( + "otel.status_code", + if outcome == "ok" { "OK" } else { "ERROR" }, + ); + if let Some(error) = error { + self.span.record("error.type", error); + } + if !self.accounting { + return; + } + record_gen_ai_usage(&self.span, &self.usage); + // Aggregate semantic counts cannot distinguish missing uncached input from zero. + self.span.record( + "switchyard.usage.input_tokens_reported", + self.usage.input_tokens.is_some(), + ); + if let Some(input) = self.usage.input_tokens { + self.span.record("switchyard.usage.input_tokens", input); + } + if let Some(total) = self.usage.total_tokens { + self.span.record("switchyard.usage.total_tokens", total); + } + if let Some(cost) = self.usage.provider_cost { + self.span.record("switchyard.usage.provider_cost_usd", cost); + } + let known = self.usage.input_tokens.is_some() + || self.usage.output_tokens.is_some() + || self.usage.cached_input_tokens().is_some() + || self.usage.cache_creation_input_tokens().is_some() + || self + .usage + .cache + .as_ref() + .is_some_and(|cache| cache.cache_creation_1h_input_tokens.is_some()) + || self.usage.total_tokens.is_some() + || self.usage.reasoning_tokens.is_some(); + let complete = self.final_usage + && self.usage.input_tokens.is_some() + && self.usage.output_tokens.is_some(); + self.span.record( + "switchyard.accounting.usage_status", + if complete { + "complete" + } else if known { + "partial" + } else { + "missing" + }, + ); + } + + pub(crate) fn wrap(self, stream: LlmResponseStream) -> LlmResponseStream { + Box::pin(AttemptStream { + inner: stream, + guard: Some(self), + dispatch: tracing::dispatcher::get_default(Clone::clone), + }) + } +} + +impl Drop for AttemptGuard { + fn drop(&mut self) { + self.finish( + if self.model_finished { + "ok" + } else { + "cancelled" + }, + None, + ); + } +} + +struct AttemptStream { + inner: LlmResponseStream, + guard: Option, + dispatch: tracing::Dispatch, +} + +impl Stream for AttemptStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let Self { + inner, + guard, + dispatch, + } = self.get_mut(); + tracing::dispatcher::with_default(dispatch, || match inner.as_mut().poll_next(cx) { + Poll::Ready(Some(item)) => { + if guard.as_mut().is_some_and(|guard| guard.observe(&item)) { + guard.take(); + } + Poll::Ready(Some(item)) + } + Poll::Ready(None) => { + if let Some(mut guard) = guard.take() { + guard.final_usage |= guard.model_finished; + guard.success(); + } + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + }) + } +} + +impl Drop for AttemptStream { + fn drop(&mut self) { + let Self { + inner, + guard, + dispatch, + } = self; + // Registry parent references can close recursively when the last child + // drops. Keep that close on the subscriber that owns the response. + tracing::dispatcher::with_default(dispatch, || { + drop(guard.take()); + drop(std::mem::replace(inner, Box::pin(futures::stream::empty()))); + }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::{StreamExt, stream}; + use serde_json::{Map, Value}; + use std::sync::{Arc, Mutex}; + use tracing::Subscriber; + use tracing::instrument::WithSubscriber; + use tracing_subscriber::{ + Layer, + layer::{Context as LayerContext, SubscriberExt}, + registry::LookupSpan, + }; + + type Records = Arc>>>; + struct Collector(Records); + + impl LookupSpan<'a>> Layer for Collector { + fn on_new_span( + &self, + attrs: &tracing::span::Attributes<'_>, + id: &tracing::Id, + ctx: LayerContext<'_, S>, + ) { + let mut fields = Map::::new(); + attrs.record( + &mut |field: &tracing::field::Field, value: &dyn std::fmt::Debug| { + let text = format!("{value:?}"); + fields.insert( + field.name().into(), + serde_json::from_str(&text).unwrap_or(Value::String(text)), + ); + }, + ); + ctx.span(id).unwrap().extensions_mut().insert(fields); + } + fn on_record( + &self, + id: &tracing::Id, + values: &tracing::span::Record<'_>, + ctx: LayerContext<'_, S>, + ) { + let span = ctx.span(id).unwrap(); + let mut extensions = span.extensions_mut(); + let fields = extensions.get_mut::>().unwrap(); + values.record( + &mut |field: &tracing::field::Field, value: &dyn std::fmt::Debug| { + let text = format!("{value:?}"); + fields.insert( + field.name().into(), + serde_json::from_str(&text).unwrap_or(Value::String(text)), + ); + }, + ); + } + fn on_close(&self, id: tracing::Id, ctx: LayerContext<'_, S>) { + let fields = ctx + .span(&id) + .unwrap() + .extensions_mut() + .remove::>() + .unwrap(); + self.0.lock().unwrap().push(fields); + } + } + + fn usage(input: Option, output: Option) -> LlmResponseChunk { + LlmResponseChunk::Usage(Usage { + input_tokens: input, + output_tokens: output, + ..Usage::default() + }) + } + + async fn observe( + events: Vec, + take: Option, + ) -> Map { + observe_results(events.into_iter().map(Ok).collect(), take).await + } + + async fn observe_results( + events: Vec>, + take: Option, + ) -> Map { + let records = Records::default(); + let subscriber = tracing_subscriber::registry().with(Collector(records.clone())); + async { + let guard = + AttemptGuard::new(&ModelId::from("model"), WireFormat::OpenAiChat, 1, 1, true); + let mut response = guard.wrap(stream::iter(events).boxed()); + if let Some(take) = take { + for _ in 0..take { + response.next().await; + } + } else { + while response.next().await.is_some() {} + } + drop(response); + } + .with_subscriber(subscriber) + .await; + let records = records.lock().unwrap(); + assert_eq!(records.len(), 1); + records[0].clone() + } + + #[tokio::test] + async fn usage_in_a_terminal_error_event_survives_regardless_of_chunk_order() { + let fields = observe( + vec![LlmResponseStreamEvent::new(vec![ + LlmResponseChunk::StreamError { + message: "private failure".into(), + }, + usage(Some(11), Some(3)), + ])], + None, + ) + .await; + assert_eq!(fields["outcome"], "error"); + assert_eq!(fields["gen_ai.usage.input_tokens"], 11); + assert_eq!(fields["gen_ai.usage.output_tokens"], 3); + assert!(!fields.contains_key("error")); + } + + #[tokio::test] + async fn cumulative_snapshots_preserve_missing_fields_and_final_corrections() { + let fields = observe( + vec![ + usage(Some(11), Some(2)).into(), + usage(None, Some(3)).into(), + usage(Some(9), Some(1)).into(), + LlmResponseChunk::MessageStop { + reason: Some("stop".into()), + } + .into(), + ], + None, + ) + .await; + assert_eq!(fields["gen_ai.usage.input_tokens"], 9); + assert_eq!(fields["gen_ai.usage.output_tokens"], 1); + assert_eq!(fields["switchyard.accounting.usage_status"], "complete"); + } + + #[tokio::test] + async fn drop_after_finish_is_success_with_partial_accounting() { + let fields = observe( + vec![ + usage(Some(11), Some(3)).into(), + LlmResponseChunk::MessageStop { + reason: Some("stop".into()), + } + .into(), + usage(Some(12), Some(4)).into(), + ], + Some(2), + ) + .await; + assert_eq!(fields["outcome"], "ok"); + assert_eq!(fields["gen_ai.usage.input_tokens"], 11); + assert_eq!(fields["switchyard.accounting.usage_status"], "partial"); + } + + #[tokio::test] + async fn missing_usage_is_distinct_from_explicit_zero() { + for (counts, status) in [(None, "missing"), (Some(0), "complete")] { + let fields = observe( + vec![ + LlmResponseChunk::MessageStop { + reason: Some("stop".into()), + } + .into(), + usage(counts, counts).into(), + ], + None, + ) + .await; + assert_eq!(fields["switchyard.accounting.usage_status"], status); + assert_eq!( + fields + .get("gen_ai.usage.input_tokens") + .and_then(Value::as_u64), + counts + ); + } + } + + #[tokio::test] + async fn trailing_usage_after_model_finish_is_not_lost_or_added_twice() { + let fields = observe( + vec![ + usage(Some(11), Some(3)).into(), + LlmResponseChunk::MessageStop { + reason: Some("stop".into()), + } + .into(), + usage(Some(12), Some(4)).into(), + usage(Some(12), Some(4)).into(), + ], + None, + ) + .await; + assert_eq!(fields["gen_ai.usage.input_tokens"], 12); + assert_eq!(fields["gen_ai.usage.output_tokens"], 4); + assert_eq!(fields["switchyard.accounting.usage_status"], "complete"); + } + + #[tokio::test] + async fn timeout_after_usage_retains_partial_consumption_and_error() { + let fields = observe_results( + vec![ + Ok(usage(Some(11), Some(3)).into()), + Err(LlmClientError::Timeout { + source: Box::new(std::io::Error::from(std::io::ErrorKind::TimedOut)), + }), + ], + None, + ) + .await; + assert_eq!(fields["gen_ai.usage.input_tokens"], 11); + assert_eq!(fields["outcome"], "error"); + assert_eq!(fields["error.type"], "timeout"); + assert_eq!(fields["switchyard.accounting.usage_status"], "partial"); + } +} diff --git a/crates/libsy-llm-client/src/client.rs b/crates/libsy-llm-client/src/client.rs index 7962f1ec5b..a7ef0ea4df 100644 --- a/crates/libsy-llm-client/src/client.rs +++ b/crates/libsy-llm-client/src/client.rs @@ -26,6 +26,7 @@ use switchyard_translation::{ use tokio::time::Instant; use tracing::Instrument; +use crate::attempt::AttemptGuard; use crate::backend::{Backend, openai_url}; use crate::error::{LlmClientError, Result}; use crate::metrics; @@ -414,7 +415,15 @@ impl TranslatingLlmClient { record_gen_ai_request(&url, provider_model_id, streaming); let result = self - .send_with_retries(&url, backend, &body, metadata, target_name, streaming) + .send_with_retries( + &url, + backend, + &body, + metadata, + target_name, + streaming, + matches!(endpoint, UpstreamEndpoint::Completion), + ) .await; // Keep quota failures request-local when forwarding auth or returning 429 to callers. if let Some(until) = unavailable_until @@ -431,6 +440,7 @@ impl TranslatingLlmClient { } // Sends the encoded body, retrying retryable failures within the backend's retry budget. + #[allow(clippy::too_many_arguments)] async fn send_with_retries( &self, url: &str, @@ -439,6 +449,7 @@ impl TranslatingLlmClient { metadata: Option<&Metadata>, target_name: &ModelId, streaming: bool, + accounting: bool, ) -> Result { let max_retries = u64::from(backend.max_retries()); let max_attempts = max_retries @@ -469,40 +480,34 @@ impl TranslatingLlmClient { { return Err(deadline_error(timeout)); } - let span = tracing::debug_span!( - target: "libsy", - "libsy.upstream_attempt", - model = %target_name, - wire_format = %backend.wire_format(), - attempt = attempt + 1 + u64::from(template_retry), - max_attempts, - retry = attempt > 0 || template_retry, - openinference.span.kind = "CHAIN", - outcome = tracing::field::Empty, - status_code = tracing::field::Empty, - will_retry = tracing::field::Empty, - retry_delay_ms = tracing::field::Empty, - ); - let mut attempt_started = false; + let request = self.build_request(url, backend, body.as_ref(), metadata)?; + let mut guard = None; let result = tokio::select! { biased; - timeout = &mut deadline => { - if attempt_started { - metrics::record_upstream_attempt(None); - } - span.record("outcome", "error"); - span.record("will_retry", false); - return Err(deadline_error(timeout)); - } + _ = &mut deadline => None, result = async { - attempt_started = true; - self.send_once(url, backend, body.as_ref(), metadata, target_name, streaming).await - }.instrument(span.clone()) => result, + let attempt_guard = guard.insert(AttemptGuard::new(self.provider_model(target_name), backend.wire_format(), attempt + 1 + u64::from(template_retry), max_attempts, accounting)); + let span = attempt_guard.span().clone(); + self.send_once(request, backend, metadata, target_name, streaming, attempt_guard).instrument(span).await + } => Some(result), + }; + let Some(result) = result else { + let timeout = expires_at + .expect("only a configured deadline can resolve") + .1; + let error = deadline_error(timeout); + if let Some(guard) = &mut guard { + metrics::record_upstream_attempt(None); + guard.span().record("will_retry", false); + guard.error(&error); + } + return Err(error); }; + let mut guard = guard.expect("a completed HTTP attempt has accounting state"); + let span = guard.span().clone(); // The retained handle updates this same attempt span with its outcome. match result { Ok(response) => { - span.record("outcome", "ok"); span.record("status_code", response.status()); span.record("will_retry", false); if attempt > 0 || template_retry { @@ -515,10 +520,25 @@ impl TranslatingLlmClient { upstream_headers, } => EncodedResponse::Streaming { status, - chunks: stream_with_deadline(chunks, expires_at), + chunks: guard.wrap(stream_with_deadline(chunks, expires_at)), upstream_headers, }, - buffered => buffered, + EncodedResponse::Buffered { + status, + body, + upstream_headers, + .. + } => { + if !accounting { + guard.success(); + } + EncodedResponse::Buffered { + status, + body, + upstream_headers, + attempt: accounting.then_some(guard), + } + } }); } Err(failure) => { @@ -528,7 +548,7 @@ impl TranslatingLlmClient { && rejects_template_options(&failure.error); let will_retry = retry_without_template || (attempt < max_retries && failure.is_retryable()); - span.record("outcome", "error"); + guard.error(&failure.error); if let Some(status) = failure.status { span.record("status_code", status.as_u16()); } @@ -549,6 +569,7 @@ impl TranslatingLlmClient { span.record("retry_delay_ms", duration_millis(delay)); // Close the attempt span before sleeping so backoff is not attempt latency. drop(span); + drop(guard); tokio::select! { biased; timeout = &mut deadline => return Err(deadline_error(timeout)), @@ -560,16 +581,14 @@ impl TranslatingLlmClient { } } - // Performs one HTTP attempt and retains the retry metadata alongside any error. - async fn send_once( + // Validate the request before opening a physical attempt. + fn build_request( &self, url: &str, backend: &Backend, body: &Value, metadata: Option<&Metadata>, - target_name: &ModelId, - streaming: bool, - ) -> std::result::Result { + ) -> Result { let client = if backend.is_forwarding_auth() { &self.forward_auth_client } else { @@ -580,9 +599,24 @@ impl TranslatingLlmClient { let builder = backend.apply_forwarded_auth(builder, metadata); let builder = apply_extra_headers(builder, backend); let builder = backend.apply_auth(builder); + builder.build().map_err(convert_reqwest_error) + } + async fn send_once( + &self, + request: reqwest::Request, + backend: &Backend, + metadata: Option<&Metadata>, + target_name: &ModelId, + streaming: bool, + guard: &mut AttemptGuard, + ) -> std::result::Result { + let client = if backend.is_forwarding_auth() { + &self.forward_auth_client + } else { + &self.client + }; let response = match async { - let request = builder.build()?; if let Some(body) = request.body().and_then(reqwest::Body::as_bytes) { crate::payload_capture::record("request", || String::from_utf8_lossy(body)); } @@ -601,20 +635,23 @@ impl TranslatingLlmClient { } }; let status = response.status(); + guard.span().record("status_code", status.as_u16()); + guard.headers(response.headers()); if status.is_success() { if streaming { let upstream_headers = response.headers().clone(); - let chunks = match prepare_response_stream(response, backend, target_name).await { - Ok(chunks) => chunks, - Err(error) => { - metrics::record_upstream_attempt(None); - return Err(AttemptFailure { - error, - status: Some(status), - retry_after: None, - }); - } - }; + let chunks = + match prepare_response_stream(response, backend, target_name, guard).await { + Ok(chunks) => chunks, + Err(error) => { + metrics::record_upstream_attempt(None); + return Err(AttemptFailure { + error, + status: Some(status), + retry_after: None, + }); + } + }; metrics::record_upstream_attempt(Some(status.as_u16())); return Ok(EncodedResponse::Streaming { status: status.as_u16(), @@ -640,6 +677,7 @@ impl TranslatingLlmClient { status: status.as_u16(), body: body.to_vec(), upstream_headers, + attempt: None, }); } @@ -655,6 +693,9 @@ impl TranslatingLlmClient { }); } }; + if let Ok(value) = serde_json::from_str::(&body) { + guard.response(&value, backend.wire_format()); + } let body = redact_forwarded_headers(body, metadata, backend.is_forwarding_auth()); metrics::record_upstream_attempt(Some(status.as_u16())); let error = @@ -681,6 +722,8 @@ impl TranslatingLlmClient { request: Request, target_name: Option<&ModelId>, ) -> Result { + tracing::Span::current().record("switchyard.accounting.unit", "summary"); + tracing::Span::current().record("switchyard.accounting.attempts", 0_u64); let Request { mut llm_request, metadata, @@ -732,16 +775,22 @@ impl TranslatingLlmClient { EncodedResponse::Buffered { body, upstream_headers, + mut attempt, .. } => { - let body = serde_json::from_slice::(&body).map_err(|error| { - LlmClientError::ResponseTranslation(format!("invalid upstream JSON: {error}")) - })?; - // Map a provider's failed generation to 502, even under HTTP 200. - // Redact forwarded credentials before returning the provider error. - let mut agg = - observed_translation::decode_response(&body, wire_format).map_err(|error| { - match error { + let result = (|| { + let body = serde_json::from_slice::(&body).map_err(|error| { + LlmClientError::ResponseTranslation(format!( + "invalid upstream JSON: {error}" + )) + })?; + if let Some(guard) = &mut attempt { + guard.response(&body, wire_format); + } + // Map a provider's failed generation to 502, even under HTTP 200. + // Redact forwarded credentials before returning the provider error. + let mut agg = observed_translation::decode_response(&body, wire_format) + .map_err(|error| match error { TranslationError::UpstreamFailure { error } => { LlmClientError::UpstreamHttp { status: StatusCode::BAD_GATEWAY, @@ -753,17 +802,24 @@ impl TranslatingLlmClient { } } error => LlmClientError::ResponseTranslation(error.to_string()), - } - })?; - // The gateway may report a buffered charge only in its response header. - if agg.usage.provider_cost.is_none() { - agg.usage.provider_cost = upstream_headers - .get("x-litellm-response-cost") - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.trim().parse::().ok()) - .filter(|usd| usd.is_finite() && *usd >= 0.0); + })?; + // The gateway may report a buffered charge only in its response header. + if agg.usage.provider_cost.is_none() { + agg.usage.provider_cost = upstream_headers + .get("x-litellm-response-cost") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.trim().parse::().ok()) + .filter(|usd| usd.is_finite() && *usd >= 0.0); + } + Ok((LlmResponse::Agg(agg), upstream_headers)) + })(); + if let Some(guard) = &mut attempt { + match &result { + Ok(_) => guard.success(), + Err(error) => guard.error(error), + } } - (LlmResponse::Agg(agg), upstream_headers) + result? } }; @@ -882,6 +938,7 @@ enum EncodedResponse { status: u16, body: Vec, upstream_headers: HeaderMap, + attempt: Option, }, Streaming { status: u16, @@ -959,6 +1016,7 @@ async fn prepare_response_stream( response: reqwest::Response, backend: &Backend, target_name: &ModelId, + guard: &mut AttemptGuard, ) -> Result { let bytes = response.bytes_stream().map(|chunk| { chunk @@ -975,6 +1033,9 @@ async fn prepare_response_stream( Some(first) => { // An in-band context overflow skips retries and advances to another candidate. if let Some(message) = first_event_overflow(&first, backend) { + if let Ok(event) = &first { + guard.observe_usage(event); + } return Err(LlmClientError::ContextWindowExceeded { model: target_name.clone(), message, diff --git a/crates/libsy-llm-client/src/lib.rs b/crates/libsy-llm-client/src/lib.rs index 3710ad150d..d216e68350 100644 --- a/crates/libsy-llm-client/src/lib.rs +++ b/crates/libsy-llm-client/src/lib.rs @@ -20,6 +20,7 @@ //! answer call from the routing outcome when needed. A host that just wants the answer does not //! have to drive the step stream itself. +mod attempt; pub mod backend; pub mod client; pub mod error; diff --git a/crates/libsy-llm-client/src/observability.rs b/crates/libsy-llm-client/src/observability.rs index e1b92bc13d..a10c0cb681 100644 --- a/crates/libsy-llm-client/src/observability.rs +++ b/crates/libsy-llm-client/src/observability.rs @@ -113,7 +113,7 @@ fn record_gen_ai_response(span: &Span, response: &AggLlmResponse) { } /// Records semantic token totals and cache-usage breakdowns on a client-call span. -fn record_gen_ai_usage(span: &Span, usage: &Usage) { +pub(crate) fn record_gen_ai_usage(span: &Span, usage: &Usage) { let cache_read = usage.cached_input_tokens(); let cache_creation = usage.cache_creation_input_tokens(); let cache_creation_1h = usage @@ -207,7 +207,7 @@ fn client_call_error_type(error: &LibsyError) -> Cow<'static, str> { } } -fn llm_client_error_type(error: &LlmClientError) -> Cow<'static, str> { +pub(crate) fn llm_client_error_type(error: &LlmClientError) -> Cow<'static, str> { match error { LlmClientError::InvalidRequest { .. } => Cow::Borrowed("invalid_request"), LlmClientError::RequestTranslation(_) => Cow::Borrowed("request_translation"), diff --git a/crates/libsy-llm-client/src/run.rs b/crates/libsy-llm-client/src/run.rs index 2ee818c851..b854fa3ed3 100644 --- a/crates/libsy-llm-client/src/run.rs +++ b/crates/libsy-llm-client/src/run.rs @@ -422,6 +422,8 @@ impl CallPhase { gen_ai.usage.cache_creation.input_tokens = tracing::field::Empty, gen_ai.usage.cache_creation_1h.input_tokens = tracing::field::Empty, gen_ai.usage.reasoning.output_tokens = tracing::field::Empty, + switchyard.accounting.unit = tracing::field::Empty, + switchyard.accounting.attempts = tracing::field::Empty, outcome = tracing::field::Empty, otel.status_code = tracing::field::Empty, error.type = tracing::field::Empty, diff --git a/crates/libsy-llm-client/tests/observability.rs b/crates/libsy-llm-client/tests/observability.rs index 509528c7f1..b1c2b8a29d 100644 --- a/crates/libsy-llm-client/tests/observability.rs +++ b/crates/libsy-llm-client/tests/observability.rs @@ -2445,6 +2445,153 @@ async fn in_flight_gauge_clears_when_a_run_is_abandoned() -> switchyard_libsy::R Ok(()) } +#[tokio::test] +async fn routing_buffer_failure_preserves_real_attempt_usage() +-> Result<(), Box> { + use switchyard_llm_client::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient}; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + + let _guard = serialize_test().lock().await; + let (store, _, _, _, _) = telemetry(); + const MODEL: &str = "obs-routing-failure-usage"; + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_raw(concat!( + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"r1\"}}\n\n", + "data: {\"type\":\"response.failed\",\"response\":{\"error\":{\"message\":\"synthetic failure\"},\"usage\":{\"input_tokens\":11,\"output_tokens\":3}}}\n\n" + ), "text/event-stream")) + .expect(1) + .mount(&server).await; + let client = Arc::new(TranslatingLlmClient::new(&[ModelConfig::new( + MODEL, + Backend::OpenAiResponses(HttpBackendConfig { + base_url: server.uri(), + api_key: None, + forward_auth: false, + extra_headers: Default::default(), + extra_body: Default::default(), + omit_body_fields: Default::default(), + reasoning_effort: None, + max_retries: 1, + failure_cooldown: Duration::ZERO, + timeout: None, + }), + None, + )])?); + let algorithm = Arc::new(RoutingCallAlgo { + name: "obs-routing-failure-usage-algo".into(), + target: MODEL.into(), + }); + let mut request = request_with_metadata("routing-failure", "routing-failure"); + request.llm_request.stream = true; + assert!(run(algorithm, client, request).await.is_err()); + let spans = store.spans(); + let span = find_span(&spans, "libsy.upstream_attempt", "model", MODEL); + assert_eq!( + span.fields + .get("gen_ai.usage.input_tokens") + .map(String::as_str), + Some("11") + ); + assert_eq!( + span.fields + .get("gen_ai.usage.output_tokens") + .map(String::as_str), + Some("3") + ); + assert_eq!( + span.fields.get("outcome").map(String::as_str), + Some("error") + ); + assert_eq!( + spans + .iter() + .filter(|s| s.name == "libsy.upstream_attempt" + && s.fields.get("model").is_some_and(|m| m == MODEL)) + .count(), + 1 + ); + Ok(()) +} + +#[tokio::test] +async fn retries_and_fallback_keep_separate_models_and_attempt_ordinals() +-> Result<(), Box> { + use switchyard_llm_client::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient}; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + + let _guard = serialize_test().lock().await; + let (store, _, _, _, _) = telemetry(); + const WEAK: &str = "obs-attempt-fallback-weak"; + const STRONG: &str = "obs-attempt-fallback-strong"; + let server = MockServer::start().await; + Mock::given(method("POST")).respond_with(|request: &wiremock::Request| { + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + if body["model"] == WEAK { + ResponseTemplate::new(503).insert_header("retry-after", "0") + .set_body_json(json!({"error":{"message":"synthetic failure"},"usage":{"prompt_tokens":11,"completion_tokens":3,"cost":0.02}})) + } else { + ResponseTemplate::new(200).set_body_json(json!({ + "id":"fallback", "model":STRONG, + "choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}], + "usage":{"prompt_tokens":7,"completion_tokens":2,"cost":0.03} + })) + } + }).expect(3).mount(&server).await; + let backend = Backend::OpenAiChat(HttpBackendConfig { + base_url: server.uri(), + api_key: None, + forward_auth: false, + extra_headers: Default::default(), + extra_body: Default::default(), + omit_body_fields: Default::default(), + reasoning_effort: None, + max_retries: 1, + failure_cooldown: Duration::ZERO, + timeout: None, + }); + let client = Arc::new(TranslatingLlmClient::new(&[ + ModelConfig::new(WEAK, backend.clone(), None), + ModelConfig::new(STRONG, backend, None), + ])?); + let algorithm = Arc::new(SingleCallAlgo { + name: "obs-attempt-fallback-algo".into(), + target_set: vec![WEAK.into(), STRONG.into()], + metadata: None, + }); + let (model, response) = run( + algorithm, + client, + request_with_metadata("fallback", "fallback"), + ) + .await?; + assert_eq!(model.as_str(), WEAK); + assert_eq!(response.served_model().unwrap().as_str(), STRONG); + let spans = store.spans(); + for (model, count, cost) in [(WEAK, 2, "0.02"), (STRONG, 1, "0.03")] { + let attempts: Vec<_> = spans + .iter() + .filter(|s| { + s.name == "libsy.upstream_attempt" + && s.fields.get("model").is_some_and(|m| m == model) + }) + .collect(); + assert_eq!(attempts.len(), count); + let mut ordinals = Vec::new(); + for attempt in attempts { + ordinals.push(attempt.fields["attempt"].parse::()?); + assert_eq!(attempt.fields["switchyard.usage.provider_cost_usd"], cost); + } + ordinals.sort(); + assert_eq!(ordinals, (1..=count).collect::>()); + assert_eq!( + find_span(&spans, "libsy.client_call", "selected_model", model).fields["switchyard.accounting.attempts"], + count.to_string() + ); + } + Ok(()) +} + /// Routing keys remain distinct while both routing-time and answer metrics aggregate by provider. #[tokio::test] async fn target_aliases_preserve_provider_model_metric_labels() -> switchyard_libsy::Result<()> { diff --git a/crates/switchyard-nemo-relay-plugin/README.md b/crates/switchyard-nemo-relay-plugin/README.md index 781c04e6de..98eeae4082 100644 --- a/crates/switchyard-nemo-relay-plugin/README.md +++ b/crates/switchyard-nemo-relay-plugin/README.md @@ -318,3 +318,29 @@ that guarantee. behavior. The plugin does not maintain a separate trusted-default target or rerun routing after an execution failure. Failures outside the shared runner, including response translation failures, are returned to Relay. + +## Validate attempt accounting + +After building and packaging the native library as described above, run the +local acceptance fixture with Relay 0.10: + +```sh +uv run --with nemo-relay==0.10.0 python \ + crates/switchyard-nemo-relay-plugin/scripts/test_attempt_accounting.py \ + --manifest /absolute/path/to/bundle/relay-plugin.toml \ + --output /tmp/switchyard-attempt-accounting +``` + +The fixture uses a loopback HTTP provider and fixed synthetic prices. It checks +failed and successful attempts, reported-charge precedence, token-price +estimates, missing usage, unknown prices, zero charges, duplicate delivery, +missing records, and workflow isolation. It writes ATOF events and separate +cost summaries. Relay also loads system plugin configuration when present. + +The [cost-summary example](examples/attempt_costs.py) selects unique attempt +records and uses Relay's codec to obtain costs. These summaries cover Switchyard +completion work. They do not change ATIF `final_metrics` or include separate +decision-model usage. + +See the [versioned accounting scope profile](../../docs/integrations/nemo_relay.md#accounting-scope-profile-version-1) +for stable fields and the limits of `coverage_complete` and `total_cost_usd`. diff --git a/crates/switchyard-nemo-relay-plugin/examples/attempt_costs.py b/crates/switchyard-nemo-relay-plugin/examples/attempt_costs.py new file mode 100644 index 0000000000..408a04937c --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/examples/attempt_costs.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Price unique Switchyard attempts from one workflow's ATOF events using Relay.""" + +import argparse +import json +import math +from pathlib import Path + +from nemo_relay import codecs + + +def summarize(events, workflow_uuid): + """Return observed spending; a missing total means coverage or pricing is incomplete. + + Configure Relay's pricing component before calling this function to enable + estimates. Provider-reported charges work without a pricing catalog. + Coverage applies only to observed instrumented completion calls, including + routing calls. It excludes separate decisions and other workflow work and + cannot detect calls whose entire telemetry is absent. + """ + scopes = {e["uuid"]: e for e in events if e.get("kind") == "scope"} + + def belongs(event): + seen = set() + while event["uuid"] not in seen: + if event["uuid"] == workflow_uuid: + return True + seen.add(event["uuid"]) + event = scopes.get(event.get("parent_uuid")) + if event is None: + return False + return False + + ends = { + e["uuid"]: e + for e in events + if e.get("kind") == "scope" and e.get("scope_category") == "end" and belongs(e) + } + calls = [e for e in scopes.values() if e.get("name") == "libsy.client_call" and belongs(e)] + attempts = [ + e for e in ends.values() + if e.get("metadata", {}).get("switchyard.accounting.unit") == "completion_attempt" + ] + coverage = bool(calls) and workflow_uuid in ends + for call in calls: + meta = ends.get(call["uuid"], {}).get("metadata", {}) + children = [e for e in attempts if e.get("parent_uuid") == call["uuid"]] + expected = meta.get("switchyard.accounting.attempts") + coverage &= ( + meta.get("switchyard.accounting.unit") == "summary" + and isinstance(expected, int) + and expected >= 0 + and expected == meta.get("switchyard.accounting.attempt_records") == len(children) + and sorted(e["metadata"]["attempt"] for e in children) == list(range(1, expected + 1)) + ) + coverage &= all(e.get("parent_uuid") in {c["uuid"] for c in calls} for e in attempts) + records = [] + for event in attempts: + meta = event["metadata"] + usage = codecs.AnthropicMessagesCodec().decode_response(event["data"]).usage or {} + cost = usage.get("cost", {}) + amount = cost.get("total") + priced = ( + cost.get("currency", "USD").upper() == "USD" + and isinstance(amount, (float, int)) and math.isfinite(amount) and amount >= 0 + ) + complete = priced and ( + cost.get("source") == "provider_reported" + or meta.get("switchyard.accounting.usage_status") == "complete" + ) + records.append({ + "uuid": event["uuid"], "model": event["data"].get("model"), + "outcome": meta.get("outcome"), + "usage_status": meta.get("switchyard.accounting.usage_status"), + "cost": cost if priced else None, "cost_complete": complete, + }) + subtotal = math.fsum(r["cost"]["total"] for r in records if r["cost"] is not None) + complete = coverage and all(r["cost_complete"] for r in records) + return { + "workflow_uuid": workflow_uuid, "attempts": records, + "coverage_complete": bool(coverage), "cost_status": "complete" if complete else "incomplete", + "known_cost_subtotal_usd": subtotal, "total_cost_usd": subtotal if complete else None, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("events", type=Path, help="ATOF JSONL containing the workflow's scope tree") + parser.add_argument("--workflow-uuid", required=True) + args = parser.parse_args() + events = [json.loads(line) for line in args.events.read_text().splitlines() if line.strip()] + print(json.dumps(summarize(events, args.workflow_uuid), indent=2)) + + +if __name__ == "__main__": + main() diff --git a/crates/switchyard-nemo-relay-plugin/scripts/test_attempt_accounting.py b/crates/switchyard-nemo-relay-plugin/scripts/test_attempt_accounting.py new file mode 100644 index 0000000000..12dc6e8e37 --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/scripts/test_attempt_accounting.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Exercise a packaged native plugin with Relay 0.10 and a synthetic provider.""" + +import argparse +import asyncio +import json +import math +import platform +import sys +import tempfile +import threading +import tomllib +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from importlib.metadata import version +from pathlib import Path + +import nemo_relay as relay +from nemo_relay import codecs, model_pricing, plugin, subscribers + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) +from attempt_costs import summarize # noqa: E402 + + +class Provider(BaseHTTPRequestHandler): + scenario = "reported" + calls = 0 + + def do_POST(self): + request = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + first = Provider.calls % 2 == 0 + Provider.calls += 1 + usage = {"prompt_tokens": 11 if first else 7, "completion_tokens": 3 if first else 2} + if self.scenario in {"reported", "missing"}: + usage["cost"] = 0.02 if first else 0.03 + if self.scenario == "zero": + usage = {"cost": 0.0} + if self.scenario == "missing" and first: + usage = {} + body = {"model": request["model"], "usage": usage} + if first: + body["error"] = {"message": "synthetic failure"} + else: + body.update(id="r1", choices=[{ + "index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop", + }]) + raw = json.dumps(body).encode() + self.send_response(503 if first else 200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(raw))) + self.send_header("Retry-After", "0") + # A body charge must override this valid header fallback. + if self.scenario == "reported": + self.send_header("x-litellm-response-cost", "1.0") + self.end_headers() + self.wfile.write(raw) + + def log_message(self, *_args): + pass + + +async def run(manifest, output): + server = ThreadingHTTPServer(("127.0.0.1", 0), Provider) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + events = [] + subscribers.register("attempt-accounting-test", lambda e: events.append(e.to_dict())) + activation = None + try: + deployment = output / "switchyard.toml" + deployment.write_text(f'''schema_version = 1 +[llm_clients.local] +format = "openai_chat" +base_url = "http://127.0.0.1:{server.server_port}/v1" +max_retries = 1 +failure_cooldown_ms = 0 +[targets.default] +id = "test/model" +llm_client = "local" +[targets.unpriced] +id = "test/unpriced" +llm_client = "local" +[routes.default] +id = "switchyard/default" +type = "passthrough" +target = "default" +[routes.unpriced] +id = "switchyard/unpriced" +type = "passthrough" +target = "unpriced" +''') + config = output / "plugins.toml" + config.write_text(f'''[[plugins.dynamic]] +manifest = {json.dumps(str(manifest))} +[plugins.dynamic.config] +switchyard_config_path = {json.dumps(str(deployment))} +[plugins.policy.overrides."nvidia.switchyard"] +attestation = "integrity_only" +''') + prices = model_pricing.ComponentSpec(model_pricing.PricingConfig(sources=[ + model_pricing.InlineSource(model_pricing.PricingCatalog(entries=[ + model_pricing.ModelPricing( + provider="anthropic", model_id="test/model", pricing_as_of="2026-11-04", + pricing_source="synthetic", rates=model_pricing.TokenPricingRates( + input_per_million=1000, output_per_million=2000, + ), + ), + ])), + ])) + activation = await plugin.initialize(plugin.PluginConfig(components=[prices]), config) + (output / "environment.json").write_text(json.dumps({ + "relay_version": version("nemo-relay"), "platform": platform.platform(), + "bundle_sha256": tomllib.loads(manifest.read_text())["integrity"]["sha256"], + }, indent=2)) + + async def forbidden(_request): + raise AssertionError("native plugin did not intercept the request") + + for scenario, expected_subtotal, expected_total in [ + ("reported", 0.05, 0.05), ("missing", 0.03, None), + ("estimated", 0.028, 0.028), ("unpriced", 0.0, None), ("zero", 0.0, 0.0), + ]: + Provider.scenario = scenario + Provider.calls = 0 + route = "unpriced" if scenario == "unpriced" else "default" + with relay.scope.scope(scenario, relay.ScopeType.Agent): + await relay.llm.execute( + "openai.chat_completions", + relay.LLMRequest({}, {"model": f"switchyard/{route}", "messages": [{"role": "user", "content": "synthetic test"}]}), + forbidden, codec=codecs.OpenAIChatCodec(), response_codec=codecs.OpenAIChatCodec(), + ) + await subscribers.flush_async() + workflow = next(e["uuid"] for e in events if e.get("name") == scenario and e.get("scope_category") == "start") + report = summarize(events, workflow) + assert Provider.calls == 2 + assert len(report["attempts"]) == 2 + assert report["attempts"][0]["outcome"] == "error" + ends = {e["uuid"]: e for e in events if e.get("scope_category") == "end"} + for record in report["attempts"]: + attempt = ends[record["uuid"]] + assert attempt["name"] == "libsy.upstream_attempt" + assert attempt["metadata"]["switchyard.accounting.schema_version"] == 1 + parent = ends[attempt["parent_uuid"]] + assert parent["name"] == "libsy.client_call" + assert parent["metadata"]["switchyard.accounting.unit"] == "summary" + assert parent["metadata"]["switchyard.accounting.schema_version"] == 1 + assert report["coverage_complete"] + assert math.isclose(report["known_cost_subtotal_usd"], expected_subtotal) + if expected_total is None: + assert report["total_cost_usd"] is None + else: + assert math.isclose(report["total_cost_usd"], expected_total) + # ATOF replay, missing records, and mixed workflows must not inflate totals. + assert summarize(events + events, workflow) == report + dropped = report["attempts"][0]["uuid"] + assert summarize([e for e in events if e["uuid"] != dropped], workflow)["total_cost_usd"] is None + if scenario == "reported": + assert all(a["cost"]["source"] == "provider_reported" for a in report["attempts"]) + if scenario == "estimated": + assert all(a["cost"]["source"] == "model_pricing" for a in report["attempts"]) + if scenario == "zero": + assert all(a["usage_status"] == "missing" for a in report["attempts"]) + (output / f"{scenario}.json").write_text(json.dumps(report, indent=2)) + print(f"{scenario}: subtotal={report['known_cost_subtotal_usd']}, total={report['total_cost_usd']}") + (output / "events.jsonl").write_text("".join(json.dumps(e) + "\n" for e in events)) + finally: + if activation is not None: + await activation.close() + subscribers.deregister("attempt-accounting-test") + server.shutdown() + server.server_close() + thread.join() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + if args.output: + args.output.mkdir(parents=True, exist_ok=True) + asyncio.run(run(args.manifest.resolve(), args.output.resolve())) + else: + with tempfile.TemporaryDirectory() as directory: + asyncio.run(run(args.manifest.resolve(), Path(directory))) + + +if __name__ == "__main__": + main() diff --git a/crates/switchyard-nemo-relay-plugin/scripts/test_attempt_costs.py b/crates/switchyard-nemo-relay-plugin/scripts/test_attempt_costs.py new file mode 100644 index 0000000000..bea86810ff --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/scripts/test_attempt_costs.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Check USD accounting with Relay's response codec, without loading a plugin.""" + +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) +from attempt_costs import summarize # noqa: E402 + + +def cost_events(*costs): + events = [ + {"kind": "scope", "uuid": "workflow", "scope_category": "end"}, + { + "kind": "scope", "uuid": "call", "parent_uuid": "workflow", + "name": "libsy.client_call", "scope_category": "end", + "metadata": { + "switchyard.accounting.unit": "summary", + "switchyard.accounting.attempts": len(costs), + "switchyard.accounting.attempt_records": len(costs), + }, + }, + ] + for ordinal, cost in enumerate(costs, 1): + events.append({ + "kind": "scope", "uuid": f"attempt-{ordinal}", "parent_uuid": "call", + "scope_category": "end", + "metadata": { + "attempt": ordinal, "switchyard.accounting.unit": "completion_attempt", + "switchyard.accounting.usage_status": "complete", "outcome": "ok", + }, + "data": { + "type": "message", "role": "assistant", "content": [], "model": "test/model", + "usage": {"input_tokens": 1, "output_tokens": 1, "cost": cost}, + }, + }) + return events + + +def report(*costs): + return summarize(cost_events(*costs), "workflow") + + +class CurrencyTests(unittest.TestCase): + def test_usd_and_missing_currency_keep_complete_total(self): + for currency in ["USD", "usd", None]: + with self.subTest(currency=currency): + cost = {"total": 0.02} + if currency is not None: + cost["currency"] = currency + result = report(cost) + self.assertEqual(result["known_cost_subtotal_usd"], 0.02) + self.assertEqual(result["total_cost_usd"], 0.02) + self.assertEqual(result["cost_status"], "complete") + + def test_non_usd_amounts_do_not_contribute_to_usd_totals(self): + for currency, amount in [("EUR", 0.02), ("JPY", 3), ("EUR", 0.0)]: + with self.subTest(currency=currency, amount=amount): + result = report({"total": amount, "currency": currency}) + self.assertTrue(result["coverage_complete"]) + self.assertEqual(result["known_cost_subtotal_usd"], 0.0) + self.assertIsNone(result["total_cost_usd"]) + self.assertEqual(result["cost_status"], "incomplete") + self.assertFalse(result["attempts"][0]["cost_complete"]) + + def test_mixed_currencies_keep_only_the_usd_subtotal(self): + result = report({"total": 0.03, "currency": "USD"}, {"total": 0.02, "currency": "EUR"}) + self.assertEqual(result["known_cost_subtotal_usd"], 0.03) + self.assertIsNone(result["total_cost_usd"]) + self.assertEqual(result["cost_status"], "incomplete") + + +class CompatibilityTests(unittest.TestCase): + def test_versioned_and_legacy_attempts_keep_completion_only_results(self): + events = cost_events({"total": 0.02}, {"total": 0.03}) + expected = summarize(events, "workflow") + for event in [events[1], events[3]]: + event["metadata"].update({ + "switchyard.accounting.schema_version": 1, "future_field": "ignored", + }) + events.append(events[3].copy()) + for name, unit in [("libsy.decision_call", "decision_call"), ("outer", "summary")]: + events.append({ + **events[2], "uuid": name, "name": name, "parent_uuid": "workflow", + "metadata": {"switchyard.accounting.unit": unit}, + }) + result = summarize(events, "workflow") + self.assertEqual(result, expected) + self.assertEqual(set(result), { + "workflow_uuid", "attempts", "coverage_complete", "cost_status", + "known_cost_subtotal_usd", "total_cost_usd", + }) + self.assertEqual(result["total_cost_usd"], 0.05) + self.assertEqual([r["uuid"] for r in result["attempts"]], ["attempt-1", "attempt-2"]) + + def test_missing_consumption_and_reported_zero_remain_distinct(self): + events = cost_events({}) + attempt = events[-1] + attempt["metadata"]["switchyard.accounting.usage_status"] = "missing" + attempt["data"].pop("usage") + missing = summarize(events, "workflow") + self.assertTrue(missing["coverage_complete"]) + self.assertEqual(missing["known_cost_subtotal_usd"], 0.0) + self.assertIsNone(missing["total_cost_usd"]) + attempt["data"]["usage"] = {"cost": { + "total": 0.0, "currency": "USD", "source": "provider_reported", + }} + zero = summarize(events, "workflow") + self.assertEqual(zero["total_cost_usd"], 0.0) + self.assertEqual(zero["attempts"][0]["usage_status"], "missing") + self.assertEqual(zero["cost_status"], "complete") + + +if __name__ == "__main__": + unittest.main() diff --git a/crates/switchyard-nemo-relay-plugin/src/attempt_tests.rs b/crates/switchyard-nemo-relay-plugin/src/attempt_tests.rs new file mode 100644 index 0000000000..b3f6e2be68 --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/src/attempt_tests.rs @@ -0,0 +1,447 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + mpsc, +}; + +use futures_util::StreamExt; +use switchyard_protocol::{Request, WireFormat, text_request}; +use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + +use crate::{config::SwitchyardConfig, runtime::SwitchyardRuntime}; + +fn runtime(server: &MockServer, format: &str, retries: u32) -> SwitchyardRuntime { + let config = serde_json::from_value::(serde_json::json!({ + "switchyard_config": { + "schema_version": 1, + "llm_clients": {"local": {"format": format, "base_url": format!("{}/v1",server.uri()), "max_retries": retries, "failure_cooldown_ms": 0}}, + "targets": {"default": {"id": "test/model", "llm_client": "local"}}, + "routes": {"default": {"id": "switchyard/default", "type": "passthrough", "target": "default"}} + } + })).unwrap(); + SwitchyardRuntime::new(config).unwrap() +} + +fn request(stream: bool) -> Request { + let mut llm_request = text_request(Some("switchyard/default".into()), "test"); + llm_request.stream = stream; + Request { + llm_request, + raw_request: None, + metadata: None, + } +} + +fn recorder(payloads: bool) -> (SpanCapture, mpsc::Receiver) { + let (tx, rx) = mpsc::channel(); + ( + SpanCapture::new(payloads.then_some(8192), move |span| tx.send(span).unwrap()), + rx, + ) +} + +fn success() -> ResponseTemplate { + ResponseTemplate::new(200).insert_header("x-litellm-response-cost", "1.0").set_body_json(serde_json::json!({ + "id":"response", "model":"test/model", "choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}], + "usage":{"prompt_tokens":7,"completion_tokens":2,"total_tokens":9,"cost":0.03} + })) +} + +#[tokio::test] +async fn retry_usage_and_reported_charges_survive_failure_without_parent_charge() { + let server = MockServer::start().await; + let calls = Arc::new(AtomicUsize::new(0)); + let seen = calls.clone(); + Mock::given(method("POST")).respond_with(move |_: &wiremock::Request| { + if seen.fetch_add(1, Ordering::SeqCst) == 0 { + ResponseTemplate::new(503).insert_header("retry-after", "0").set_body_json(serde_json::json!({ + "error":{"message":"failed"}, "usage":{"prompt_tokens":11,"completion_tokens":3,"cost":0.02} + })) + } else { success() } + }).mount(&server).await; + let runtime = runtime(&server, "openai_chat", 1); + let (capture, rx) = recorder(false); + let result = capture + .capture(runtime.execute_buffered(WireFormat::OpenAiChat, request(false))) + .await; + assert!(result.result.is_ok()); + let call = rx.try_recv().unwrap(); + assert_eq!(call.name, CLIENT_CALL); + assert_eq!(call.children.len(), 2); + assert_eq!(call.fields["switchyard.accounting.attempts"], 2); + let (kind, metadata, _, output) = call.projection(); + assert!(matches!(kind, ScopeType::Custom)); + assert!(output.is_none()); + assert_eq!(metadata["switchyard.accounting.schema_version"], 1); + assert_eq!(metadata["switchyard.accounting.unit"], "summary"); + assert_eq!(metadata["switchyard.accounting.attempt_records"], 2); + for (attempt, outcome, input, cost) in [ + (&call.children[0], "error", 11, 0.02), + (&call.children[1], "ok", 7, 0.03), + ] { + let (kind, metadata, input_data, output) = attempt.projection(); + assert!(matches!(kind, ScopeType::Llm)); + assert_eq!(attempt.name, UPSTREAM_ATTEMPT); + assert_eq!(metadata["switchyard.accounting.schema_version"], 1); + assert_eq!(metadata["switchyard.accounting.unit"], "completion_attempt"); + assert_eq!(metadata["outcome"], outcome); + assert_eq!(metadata["switchyard.accounting.usage_status"], "complete"); + assert_eq!(input_data.unwrap()["model"], "test/model"); + let output = output.unwrap(); + assert_eq!(output["model"], "test/model"); + assert_eq!(output["usage"]["input_tokens"], input); + assert_eq!(output["usage"]["cost"]["total"], cost); + assert_eq!(output["usage"]["cost"]["currency"], "USD"); + assert_eq!(output["usage"]["cost"]["source"], "provider_reported"); + } + assert_eq!( + call.children[1].fields["switchyard.accounting.charge_source"], + "provider" + ); + assert_eq!(calls.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn cache_usage_preserves_missing_and_zero_uncached_input() { + for input in [None, Some(0), Some(7)] { + let server = MockServer::start().await; + let mut usage = serde_json::json!({"cache_read_input_tokens":10,"output_tokens":2}); + if let Some(input) = input { + usage["input_tokens"] = Json::from(input); + } + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id":"response", "type":"message", "role":"assistant", "model":"test/model", + "content":[{"type":"text","text":"ok"}], "stop_reason":"end_turn", "usage":usage + }))) + .expect(1) + .mount(&server) + .await; + let runtime = runtime(&server, "anthropic_messages", 0); + let (capture, rx) = recorder(false); + assert!( + capture + .capture(runtime.execute_buffered(WireFormat::AnthropicMessages, request(false))) + .await + .result + .is_ok() + ); + let call = rx.try_recv().unwrap(); + let (_, metadata, _, output) = call.children[0].projection(); + assert_eq!( + metadata["switchyard.usage.input_tokens_reported"], + input.is_some() + ); + assert_eq!( + metadata + .get("switchyard.usage.input_tokens") + .and_then(Json::as_u64), + input + ); + let usage = &output.as_ref().unwrap()["usage"]; + assert_eq!(usage.get("input_tokens").and_then(Json::as_u64), input); + assert_eq!(usage["cache_read_input_tokens"], 10); + assert_eq!( + metadata["switchyard.accounting.usage_status"], + if input.is_some() { + "complete" + } else { + "partial" + } + ); + } +} + +#[tokio::test] +async fn missing_retry_usage_is_not_zero() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(503).insert_header("retry-after", "0")) + .mount(&server) + .await; + let runtime = runtime(&server, "openai_chat", 1); + let (capture, rx) = recorder(false); + assert!( + capture + .capture(runtime.execute_buffered(WireFormat::OpenAiChat, request(false))) + .await + .result + .is_err() + ); + let call = rx.try_recv().unwrap(); + assert_eq!(call.children.len(), 2); + for attempt in &call.children { + let (_, fields, _, output) = attempt.projection(); + assert_eq!(fields["switchyard.accounting.usage_status"], "missing"); + assert!(output.unwrap().get("usage").is_none()); + } +} + +#[tokio::test] +async fn buffered_failed_generation_keeps_usage_without_retrying_content_failure() { + let server = MockServer::start().await; + Mock::given(method("POST")).respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id":"failed", "model":"test/model", "status":"failed", "error":{"message":"generation failed"}, + "usage":{"input_tokens":11,"output_tokens":3,"total_tokens":14} + }))).expect(1).mount(&server).await; + let runtime = runtime(&server, "openai_responses", 1); + let (capture, rx) = recorder(false); + assert!( + capture + .capture(runtime.execute_buffered(WireFormat::OpenAiResponses, request(false))) + .await + .result + .is_err() + ); + let call = rx.try_recv().unwrap(); + let [attempt] = call.children.as_slice() else { + panic!("expected one HTTP attempt") + }; + assert_eq!(attempt.fields["outcome"], "error"); + assert_eq!(attempt.fields["gen_ai.usage.input_tokens"], 11); + assert_eq!( + attempt.fields["switchyard.accounting.usage_status"], + "complete" + ); +} + +#[tokio::test] +async fn failed_stream_terminal_frame_keeps_usage_and_error() { + let server = MockServer::start().await; + let body = concat!( + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"r1\",\"model\":\"test/model\"}}\n\n", + "data: {\"type\":\"response.failed\",\"response\":{\"error\":{\"message\":\"failed\"},\"usage\":{\"input_tokens\":11,\"output_tokens\":3}}}\n\n" + ); + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(body), + ) + .expect(1) + .mount(&server) + .await; + let runtime = runtime(&server, "openai_responses", 1); + let (capture, rx) = recorder(false); + capture + .capture(async { + let mut events = runtime + .execute_stream(WireFormat::OpenAiResponses, request(true), Arc::new(|_| {})) + .await + .result + .unwrap(); + while events.next().await.is_some() {} + }) + .await; + let call = rx.try_recv().unwrap(); + let [attempt] = call.children.as_slice() else { + panic!("expected one HTTP attempt") + }; + assert_eq!(attempt.fields["outcome"], "error"); + assert_eq!(attempt.fields["gen_ai.usage.output_tokens"], 3); + assert_eq!( + attempt.fields["switchyard.accounting.usage_status"], + "complete" + ); +} + +#[tokio::test] +async fn abandoning_stream_preserves_partial_usage() { + let server = MockServer::start().await; + let body = concat!( + "data: {\"id\":\"r1\",\"model\":\"test/model\",\"choices\":[],\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":3}}\n\n", + "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\n" + ); + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(body), + ) + .mount(&server) + .await; + let runtime = runtime(&server, "openai_chat", 0); + let (capture, rx) = recorder(false); + capture + .capture(async { + let mut events = runtime + .execute_stream(WireFormat::OpenAiChat, request(true), Arc::new(|_| {})) + .await + .result + .unwrap(); + while let Some(event) = events.next().await { + if event.unwrap().to_string().contains("partial") { + break; + } + } + drop(events); + }) + .await; + let call = rx.try_recv().unwrap(); + let attempt = &call.children[0]; + assert_eq!(attempt.fields["outcome"], "cancelled"); + assert_eq!(attempt.fields["gen_ai.usage.input_tokens"], 11); + assert_eq!( + attempt.fields["switchyard.accounting.usage_status"], + "partial" + ); +} + +#[test] +fn debug_payloads_cannot_replace_usage_and_summary_is_not_priceable() { + let fields = serde_json::json!({"gen_ai.request.model":"model", "gen_ai.usage.input_tokens":12, "gen_ai.usage.cache_read.input_tokens":5, "gen_ai.usage.output_tokens":0, "switchyard.accounting.unit":"completion_attempt"}).as_object().unwrap().clone(); + let mut attempt = CapturedSpan { + name: UPSTREAM_ATTEMPT, + fields, + payloads: Payloads { + request: Some(serde_json::json!({"body":"{}","truncated":false})), + response: Some(serde_json::json!({"body":null,"truncated":true})), + }, + started_at: SystemTime::now(), + ended_at: SystemTime::now(), + children: Vec::new(), + }; + let (_, _, input, output) = attempt.projection(); + assert_eq!(input.unwrap()["model"], "model"); + let output = output.unwrap(); + assert_eq!(output["usage"]["input_tokens"], 7); + assert_eq!(output["usage"]["output_tokens"], 0); + assert_eq!(output["truncated"], true); + attempt.fields.insert( + "switchyard.accounting.unit".into(), + Json::String("summary".into()), + ); + let (kind, _, _, output) = attempt.projection(); + assert!(matches!(kind, ScopeType::Custom)); + assert!(output.is_none()); +} + +#[tokio::test] +async fn header_only_charge_is_preserved_on_an_http_error() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(400).insert_header("x-litellm-response-cost", "0")) + .expect(1) + .mount(&server) + .await; + let runtime = runtime(&server, "openai_chat", 0); + let (capture, rx) = recorder(false); + assert!( + capture + .capture(runtime.execute_buffered(WireFormat::OpenAiChat, request(false))) + .await + .result + .is_err() + ); + let call = rx.try_recv().unwrap(); + let [attempt] = call.children.as_slice() else { + panic!("expected one attempt") + }; + let (_, meta, _, output) = attempt.projection(); + assert_eq!(meta["outcome"], "error"); + assert_eq!(meta["switchyard.accounting.usage_status"], "missing"); + assert_eq!(meta["switchyard.accounting.charge_source"], "header"); + assert_eq!(output.unwrap()["usage"]["cost"]["total"], 0.0); +} + +#[tokio::test] +async fn expired_deadline_does_not_create_an_attempt() { + let server = MockServer::start().await; + use switchyard_llm_client::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient}; + use tracing::Instrument; + let client = TranslatingLlmClient::new(&[ModelConfig::new( + "test/model", + Backend::OpenAiChat(HttpBackendConfig { + base_url: server.uri(), + api_key: None, + forward_auth: false, + extra_headers: Default::default(), + extra_body: Default::default(), + omit_body_fields: Default::default(), + reasoning_effort: None, + max_retries: 0, + failure_cooldown: std::time::Duration::ZERO, + timeout: Some(std::time::Duration::ZERO), + }), + None, + )]) + .unwrap(); + let (capture, rx) = recorder(false); + let result = capture + .capture(async { + client + .call_rewrite_model(request(false), Some(&"test/model".into())) + .instrument(tracing::debug_span!(target: "libsy", "libsy.client_call", + switchyard.accounting.unit = tracing::field::Empty, + switchyard.accounting.attempts = tracing::field::Empty)) + .await + }) + .await; + assert!(result.is_err()); + let call = rx.try_recv().unwrap(); + assert!(call.children.is_empty()); + assert_eq!(call.fields["switchyard.accounting.attempts"], 0); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn concurrent_calls_keep_their_own_attempts() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(success().set_delay(std::time::Duration::from_millis(10))) + .expect(2) + .mount(&server) + .await; + let runtime = runtime(&server, "openai_chat", 0); + let (capture, rx) = recorder(false); + let first = capture.capture(runtime.execute_buffered(WireFormat::OpenAiChat, request(false))); + let second = capture.capture(runtime.execute_buffered(WireFormat::OpenAiChat, request(false))); + let (first, second) = tokio::join!(first, second); + assert!(first.result.is_ok()); + assert!(second.result.is_ok()); + let calls: Vec<_> = rx.try_iter().collect(); + assert_eq!(calls.len(), 2); + for call in calls { + assert_eq!(call.children.len(), 1); + assert_eq!(call.fields["switchyard.accounting.attempts"], 1); + assert_eq!(call.children[0].fields["gen_ai.usage.input_tokens"], 7); + } +} + +#[tokio::test] +async fn prefetched_context_rejection_keeps_usage_and_specific_error_type() { + let server = MockServer::start().await; + let body = "data: {\"type\":\"response.failed\",\"response\":{\"error\":{\"message\":\"maximum context length exceeded\"},\"usage\":{\"input_tokens\":11,\"output_tokens\":0}}}\n\n"; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_raw(body, "text/event-stream")) + .expect(1) + .mount(&server) + .await; + let runtime = runtime(&server, "openai_responses", 1); + let (capture, rx) = recorder(false); + assert!( + capture + .capture(runtime.execute_stream( + WireFormat::OpenAiResponses, + request(true), + Arc::new(|_| {}) + )) + .await + .result + .is_err() + ); + let call = rx.try_recv().unwrap(); + let [attempt] = call.children.as_slice() else { + panic!("expected one attempt") + }; + assert_eq!(attempt.fields["gen_ai.usage.input_tokens"], 11); + assert_eq!(attempt.fields["error.type"], "context_window_exceeded"); + assert_eq!( + attempt.fields["switchyard.accounting.usage_status"], + "complete" + ); +} diff --git a/crates/switchyard-nemo-relay-plugin/src/span_capture.rs b/crates/switchyard-nemo-relay-plugin/src/span_capture.rs index 59db88c4bc..767c99b0b3 100644 --- a/crates/switchyard-nemo-relay-plugin/src/span_capture.rs +++ b/crates/switchyard-nemo-relay-plugin/src/span_capture.rs @@ -200,28 +200,10 @@ impl CapturedSpan { } pub(crate) fn emit(self, runtime: &PluginRuntime) -> Result<(), String> { - let mut metadata = Json::Object(self.fields); - let output = self.payloads.response.or_else(|| { - if self.name == CLIENT_CALL { - response_data(&metadata) - } else { - None - } - }); - if let Some(session_id) = metadata.get("gen_ai.conversation.id").cloned() { - metadata["session_id"] = session_id; - } - if let Some(status) = span_status(&metadata) { - metadata["otel.status_code"] = status["otel.status_code"].clone(); - } - let input = self.payloads.request.or_else(|| { - metadata - .get("gen_ai.request.model") - .map(|model| serde_json::json!({"model": model})) - }); + let (scope_type, metadata, input, output) = self.projection(); let mut scope = runtime.scope_at( self.name, - relay_scope_type(self.name), + scope_type, None, Some(&metadata), input.as_ref(), @@ -234,6 +216,72 @@ impl CapturedSpan { let close_result = scope.close_at(output.as_ref(), Some(&metadata), self.ended_at); child_result.and(close_result) } + + fn projection(&self) -> (ScopeType, Json, Option, Option) { + let mut metadata = Json::Object(self.fields.clone()); + if let Some(session) = metadata.get("gen_ai.conversation.id").cloned() { + metadata["session_id"] = session; + } + if let Some(status) = span_status(&metadata) { + metadata["otel.status_code"] = status["otel.status_code"].clone(); + } + let unit = metadata + .get("switchyard.accounting.unit") + .and_then(Json::as_str); + let summary = unit == Some("summary"); + let attempt = unit == Some("completion_attempt"); + if summary || attempt { + metadata["switchyard.accounting.schema_version"] = Json::from(1); + } + let scope_type = if summary { + ScopeType::Custom + } else if attempt { + ScopeType::Llm + } else { + relay_scope_type(self.name) + }; + let input = request_data(&metadata, self.payloads.request.as_ref()); + let output = if summary { + metadata["switchyard.accounting.attempt_records"] = Json::from( + self.children + .iter() + .filter(|child| { + child + .fields + .get("switchyard.accounting.unit") + .and_then(Json::as_str) + == Some("completion_attempt") + }) + .count(), + ); + None + } else if attempt || self.name == CLIENT_CALL { + // Debug payloads supplement accounting, rather than replacing its usage. + let mut output = response_data(&metadata); + if let Some(payload) = &self.payloads.response { + let value = output.get_or_insert_with(|| serde_json::json!({})); + if let (Some(target), Some(source)) = (value.as_object_mut(), payload.as_object()) { + target.extend(source.clone()); + } + } + output + } else { + self.payloads.response.clone() + }; + (scope_type, metadata, input, output) + } +} + +// Preserve the requested model alongside an optional debug payload. +fn request_data(fields: &Json, payload: Option<&Json>) -> Option { + let mut input = payload.cloned().unwrap_or_else(|| serde_json::json!({})); + if let Some(model) = fields + .get("gen_ai.request.model") + .or_else(|| fields.get("selected_model")) + { + input["model"] = model.clone(); + } + (!input.as_object()?.is_empty()).then_some(input) } fn redact_captured_json(value: Json, redactor: &ProviderKeyRedactor, depth: usize) -> Option { @@ -291,9 +339,17 @@ fn response_data(fields: &Json) -> Option { output.insert(target.into(), value.clone()); } } + if !output.contains_key("model") + && let Some(model) = fields + .get("gen_ai.request.model") + .or_else(|| fields.get("selected_model")) + { + output.insert("model".into(), model.clone()); + } let mut usage = Map::new(); for (source, target) in [ ("gen_ai.usage.output_tokens", "output_tokens"), + ("gen_ai.usage.reasoning.output_tokens", "reasoning_tokens"), ( "gen_ai.usage.cache_read.input_tokens", "cache_read_input_tokens", @@ -307,28 +363,43 @@ fn response_data(fields: &Json) -> Option { usage.insert(target.into(), value.clone()); } } + if let Some(cost) = fields.get("switchyard.usage.provider_cost_usd") { + usage.insert( + "cost".into(), + serde_json::json!({"total": cost, "currency": "USD", "source": "provider_reported"}), + ); + } if let Some(value) = fields.get("gen_ai.usage.cache_creation_1h.input_tokens") { usage.insert( "cache_creation".into(), serde_json::json!({"ephemeral_1h_input_tokens": value}), ); } - if let Some(input) = fields - .get("gen_ai.usage.input_tokens") - .and_then(Json::as_u64) + let input = match fields + .get("switchyard.usage.input_tokens_reported") + .and_then(Json::as_bool) { - let cached = usage - .get("cache_read_input_tokens") + Some(true) => fields + .get("switchyard.usage.input_tokens") + .and_then(Json::as_u64), + Some(false) => None, + None => fields + .get("gen_ai.usage.input_tokens") .and_then(Json::as_u64) - .unwrap_or(0); - let created = usage - .get("cache_creation_input_tokens") - .and_then(Json::as_u64) - .unwrap_or(0); - usage.insert( - "input_tokens".into(), - input.saturating_sub(cached).saturating_sub(created).into(), - ); + .map(|input| { + let cached = usage + .get("cache_read_input_tokens") + .and_then(Json::as_u64) + .unwrap_or(0); + let created = usage + .get("cache_creation_input_tokens") + .and_then(Json::as_u64) + .unwrap_or(0); + input.saturating_sub(cached).saturating_sub(created) + }), + }; + if let Some(input) = input { + usage.insert("input_tokens".into(), input.into()); } if !usage.is_empty() { output.insert("usage".into(), Json::Object(usage)); @@ -336,8 +407,8 @@ fn response_data(fields: &Json) -> Option { if output.is_empty() { return None; } - // Relay decodes scope output with provider codecs. A content-free Messages - // response preserves separate cache-read and cache-creation token counts. + // A consumer can decode this content-free Messages response with Relay's + // codec, preserving separate cache-read and cache-creation token counts. output.insert("type".into(), Json::String("message".into())); output.insert("role".into(), Json::String("assistant".into())); output.insert("content".into(), Json::Array(Vec::new())); @@ -521,7 +592,15 @@ fn relay_scope_type(name: &str) -> ScopeType { fn is_safe_field(name: &str) -> bool { matches!( name, - "algorithm" + "switchyard.accounting.unit" + | "switchyard.accounting.attempts" + | "switchyard.accounting.usage_status" + | "switchyard.accounting.charge_source" + | "switchyard.usage.input_tokens" + | "switchyard.usage.input_tokens_reported" + | "switchyard.usage.total_tokens" + | "switchyard.usage.provider_cost_usd" + | "algorithm" | "switchyard.candidate" | "switchyard.candidate_count" | "switchyard.call_phase" @@ -1002,3 +1081,7 @@ mod tests { assert_eq!(attempt.fields["outcome"], "cancelled"); } } + +#[cfg(test)] +#[path = "attempt_tests.rs"] +mod attempt_tests; diff --git a/crates/switchyard-translation/src/codecs/anthropic/buffered.rs b/crates/switchyard-translation/src/codecs/anthropic/buffered.rs index 5639f1765c..455288afa7 100644 --- a/crates/switchyard-translation/src/codecs/anthropic/buffered.rs +++ b/crates/switchyard-translation/src/codecs/anthropic/buffered.rs @@ -1303,7 +1303,7 @@ fn encode_anthropic_tool_choice(choice: &ToolChoice) -> Value { } // Normalizes Anthropic usage fields. -fn decode_anthropic_usage(value: Option<&Value>) -> Usage { +pub(crate) fn decode_anthropic_usage(value: Option<&Value>) -> Usage { let Some(value) = value.and_then(Value::as_object) else { return Usage::default(); }; diff --git a/crates/switchyard-translation/src/codecs/anthropic/mod.rs b/crates/switchyard-translation/src/codecs/anthropic/mod.rs index d990cd988c..61701d33ff 100644 --- a/crates/switchyard-translation/src/codecs/anthropic/mod.rs +++ b/crates/switchyard-translation/src/codecs/anthropic/mod.rs @@ -8,3 +8,5 @@ mod stream; pub use buffered::AnthropicMessagesCodec; pub use stream::AnthropicMessagesStreamCodec; + +pub(crate) use buffered::decode_anthropic_usage; diff --git a/crates/switchyard-translation/src/codecs/mod.rs b/crates/switchyard-translation/src/codecs/mod.rs index 3b36a0bb82..96babe04bb 100644 --- a/crates/switchyard-translation/src/codecs/mod.rs +++ b/crates/switchyard-translation/src/codecs/mod.rs @@ -68,3 +68,16 @@ pub trait FormatCodec: Send + Sync { policy: &TranslationPolicy, ) -> Result; } + +/// Extracts reported usage independently of generation success or content validity. +/// +/// Only the provider's top-level usage object is inspected. Missing or invalid +/// counters remain unknown; the response is not accepted as a successful result. +pub fn response_usage(body: &Value, format: crate::WireFormat) -> crate::Usage { + let usage = body.get("usage"); + match format { + crate::WireFormat::OpenAiChat => openai_chat::decode_openai_usage(usage), + crate::WireFormat::OpenAiResponses => responses::decode_responses_usage(usage), + crate::WireFormat::AnthropicMessages => anthropic::decode_anthropic_usage(usage), + } +} diff --git a/crates/switchyard-translation/src/codecs/openai_chat/mod.rs b/crates/switchyard-translation/src/codecs/openai_chat/mod.rs index 766f312c6f..590fb45749 100644 --- a/crates/switchyard-translation/src/codecs/openai_chat/mod.rs +++ b/crates/switchyard-translation/src/codecs/openai_chat/mod.rs @@ -10,3 +10,5 @@ pub use buffered::OpenAiChatCodec; pub use stream::OpenAiChatStreamCodec; pub(crate) use buffered::{decode_file_source, decode_image_source}; + +pub(crate) use buffered::decode_openai_usage; diff --git a/crates/switchyard-translation/src/codecs/responses/buffered.rs b/crates/switchyard-translation/src/codecs/responses/buffered.rs index 2a0acc6ca8..b70a18eca3 100644 --- a/crates/switchyard-translation/src/codecs/responses/buffered.rs +++ b/crates/switchyard-translation/src/codecs/responses/buffered.rs @@ -2044,7 +2044,7 @@ fn json_string_python_style(value: &Value) -> String { } // Normalizes Responses usage fields. -fn decode_responses_usage(value: Option<&Value>) -> Usage { +pub(crate) fn decode_responses_usage(value: Option<&Value>) -> Usage { let Some(value) = value.and_then(Value::as_object) else { return Usage::default(); }; diff --git a/crates/switchyard-translation/src/codecs/responses/mod.rs b/crates/switchyard-translation/src/codecs/responses/mod.rs index 8f18e26fb0..8094b249ee 100644 --- a/crates/switchyard-translation/src/codecs/responses/mod.rs +++ b/crates/switchyard-translation/src/codecs/responses/mod.rs @@ -76,3 +76,5 @@ fn decode_anthropic_thinking(payload: &str) -> Option<(String, String)> { pub fn is_anthropic_thinking_envelope(encrypted_content: &str) -> bool { encrypted_content.starts_with(ANTHROPIC_THINKING_PREFIX) } + +pub(crate) use buffered::decode_responses_usage; diff --git a/crates/switchyard-translation/src/codecs/responses/stream.rs b/crates/switchyard-translation/src/codecs/responses/stream.rs index d2c49754d6..4727108a9f 100644 --- a/crates/switchyard-translation/src/codecs/responses/stream.rs +++ b/crates/switchyard-translation/src/codecs/responses/stream.rs @@ -318,24 +318,26 @@ fn decode_responses_stream( out.push(LlmResponseChunk::MessageStop { reason }); out } - // Carries the Anthropic spelling because every encoder already maps it. - Some("response.incomplete") => { + Some(kind @ ("response.incomplete" | "response.failed")) => { let mut out = Vec::new(); decode_terminal_usage(state, event, &mut out); - out.push(LlmResponseChunk::MessageStop { - reason: Some("max_tokens".to_string()), + out.push(if kind == "response.incomplete" { + LlmResponseChunk::MessageStop { + reason: Some("max_tokens".into()), + } + } else { + LlmResponseChunk::StreamError { + message: event + .get("response") + .and_then(|response| response.get("error")) + .and_then(|error| error.get("message")) + .and_then(Value::as_str) + .unwrap_or("unknown Responses stream error") + .to_string(), + } }); out } - Some("response.failed") => vec![LlmResponseChunk::StreamError { - message: event - .get("response") - .and_then(|response| response.get("error")) - .and_then(|error| error.get("message")) - .and_then(Value::as_str) - .unwrap_or("unknown Responses stream error") - .to_string(), - }], Some("error") => vec![LlmResponseChunk::StreamError { message: event .get("message") @@ -1214,7 +1216,11 @@ fn responses_usage(usage: &serde_json::Map) -> Usage { total_tokens: usage .get("total_tokens") .and_then(Value::as_u64) - .or_else(|| Some(aggregate_input_tokens.unwrap_or(0) + output_tokens.unwrap_or(0))), + .or_else(|| { + aggregate_input_tokens + .zip(output_tokens) + .map(|(input, output)| input + output) + }), reasoning_tokens: usage .get("output_tokens_details") .and_then(|details| details.get("reasoning_tokens")) diff --git a/crates/switchyard-translation/src/lib.rs b/crates/switchyard-translation/src/lib.rs index b221749cdf..2db3fdb34e 100644 --- a/crates/switchyard-translation/src/lib.rs +++ b/crates/switchyard-translation/src/lib.rs @@ -17,6 +17,7 @@ mod sse; pub mod stream; pub mod util; +pub use codecs::response_usage; pub use codecs::responses::is_anthropic_thinking_envelope; pub use switchyard_protocol::stream::{ LlmResponseChunk, LlmResponseStream, LlmResponseStreamEvent, LlmStreamError, diff --git a/crates/switchyard-translation/tests/failure_usage.rs b/crates/switchyard-translation/tests/failure_usage.rs new file mode 100644 index 0000000000..afd76686c7 --- /dev/null +++ b/crates/switchyard-translation/tests/failure_usage.rs @@ -0,0 +1,55 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use futures::{StreamExt, stream}; +use serde_json::json; +use switchyard_translation::{LlmResponseChunk, WireFormat, decode_stream, response_usage}; + +#[test] +fn failure_usage_reuses_provider_normalization_and_preserves_unknowns() { + let body = json!({"status":"failed", "error":{"message":"failed"}, "usage":{"input_tokens":20,"input_tokens_details":{"cached_tokens":3,"cache_write_tokens":2},"output_tokens":4,"cost":0.0}}); + let usage = response_usage(&body, WireFormat::OpenAiResponses); + assert_eq!(usage.input_tokens, Some(15)); + assert_eq!(usage.cached_input_tokens(), Some(3)); + assert_eq!(usage.cache_creation_input_tokens(), Some(2)); + assert_eq!(usage.total_tokens, Some(24)); + assert_eq!(usage.provider_cost, Some(0.0)); + assert_eq!( + response_usage( + &json!({"error":{"usage":{"input_tokens":10}}}), + WireFormat::OpenAiResponses + ), + Default::default() + ); + assert_eq!( + response_usage(&json!({"usage":{}}), WireFormat::AnthropicMessages), + Default::default() + ); +} + +#[test] +fn responses_terminal_failures_preserve_usage_and_original_event() { + futures::executor::block_on(async { + for kind in ["response.failed", "response.incomplete"] { + let raw = json!({"type":kind,"response":{"id":"r1","status":kind.strip_prefix("response.").unwrap(),"error":{"message":"failed"},"usage":{"input_tokens":11,"output_tokens":3}}}); + let bytes = format!("data: {raw}\n\n").into_bytes(); + let mut events = decode_stream( + stream::iter(vec![Ok(bytes)]).boxed(), + WireFormat::OpenAiResponses, + ) + .unwrap(); + let event = events.next().await.unwrap().unwrap(); + assert_eq!(event.preservation().unwrap().raw(), &raw); + let [LlmResponseChunk::Usage(usage), terminal] = event.normalized() else { + panic!("usage must precede terminal bookkeeping: {event:?}") + }; + assert_eq!(usage.input_tokens, Some(11)); + assert_eq!(usage.output_tokens, Some(3)); + assert!(matches!( + (kind, terminal), + ("response.failed", LlmResponseChunk::StreamError { .. }) + | ("response.incomplete", LlmResponseChunk::MessageStop { .. }) + )); + } + }); +} diff --git a/docs/integrations/nemo_relay.md b/docs/integrations/nemo_relay.md index 68be84f678..43b3668809 100644 --- a/docs/integrations/nemo_relay.md +++ b/docs/integrations/nemo_relay.md @@ -58,34 +58,88 @@ A routed request can spend tokens choosing a model and generating the answer. Some routes call a judge or classifier first. The selected model, or a fallback, then answers the request. -Relay records the request and answer seen by the agent. It can record a -provider-reported cost or estimate one when model pricing is configured. -Switchyard separately records model calls, available token usage and timing, -fallbacks, and failures. +Switchyard exports one `libsy.upstream_attempt` scope for each dispatched +completion attempt, including retries, routing calls, and fallback candidates. +Each scope retains available usage and supported provider charges when the +attempt fails or its stream is dropped. Local validation failures and deadlines +that expire before dispatch create no attempt. + +Use these scopes as the accounting source for Switchyard completion work: + +1. Select end events with `switchyard.accounting.unit = "completion_attempt"` + in the workflow's scope tree. Count each event UUID once. +2. Decode each scope's content-free Messages-shaped output with Relay's + `AnthropicMessagesCodec`. Relay uses reported charges first and configured + model prices otherwise. Price each attempt with its own model. +3. Exclude `libsy.client_call` summaries, token metrics, and Relay's outer answer + from this total. They describe work already represented by the attempts. +4. Check each summary's expected `switchyard.accounting.attempts` against + `switchyard.accounting.attempt_records` and the received attempt scopes. + Missing records leave coverage incomplete. + +The [ATOF cost-summary example](../../crates/switchyard-nemo-relay-plugin/examples/attempt_costs.py) +implements this selection and uses Relay's existing codec and pricing resolver. +Import `summarize(events, workflow_uuid)` into a host with Relay model pricing +configured to obtain estimates. For recorded provider charges, it also works +as a command: + +```sh +python crates/switchyard-nemo-relay-plugin/examples/attempt_costs.py events.jsonl \ + --workflow-uuid +``` -The integration does not calculate a combined cost or savings figure. Relay can -price the answer returned to the agent, while Switchyard reports internal -routing usage separately. To estimate the observed cost for a workload, use: +The result separates `known_cost_subtotal_usd` from nullable `total_cost_usd`. +For example, a failed `$0.02` attempt followed by a successful `$0.03` attempt +has a `$0.05` total. If the failed attempt has unknown consumption, the known +subtotal is `$0.03` and the complete total is absent. A reported `$0` charge is +valid even when token usage is missing. -```text -Relay answer cost + cost of Switchyard tokens where call_role = "routing" -``` +`switchyard.accounting.usage_status` is `complete`, `partial`, or `missing`. +Cumulative usage snapshots replace known fields; repeated snapshots are not +added together. Cache and reasoning fields can overlap other token counts. +Token coverage and charge coverage are separate: a known provider charge can +cover an attempt whose token usage is incomplete. -Apply the price of each `target_model` to those routing tokens. Do not add -`call_role = "answer"` again because Relay already counted that answer. Do not -sum every `token_type`; `total` is a rollup, and cache or reasoning values may -already be included in broader counts. +#### Accounting Scope Profile, Version 1 -Relay does not copy request identity into exported metric attributes. Use this -calculation for a workload or time window, not to reconcile one trace. +Completion-attempt and summary scopes set the integer metadata field +`switchyard.accounting.schema_version` to `1`. This versions the scope +representation independently of Relay's worker protocol and mark `DataSchema`. +Consumers should tolerate additional fields and unknown values. Removing or +renaming fields, changing their types, or changing their meaning requires a +new profile version. -Relay does not include a model-price catalog. Follow its -[model-pricing guide](https://docs.nvidia.com/nemo/relay/configure-plugins/model-pricing) -to supply and validate model rates. Run the same representative workload once -with a fixed model and once with routing. Compare cost, latency, fallback rate, -and task success. Missing usage, failed attempts, internal HTTP retries, and -incomplete streams can leave some cost unknown. An absent value means unknown, -not zero. +| Record | Scope type and accounting unit | Identity and output | +| --- | --- | --- | +| `libsy.client_call` | Custom; `summary` | Parent of its attempts; no priceable response output | +| `libsy.upstream_attempt` | LLM; `completion_attempt` | Its scope UUID identifies one attempt; `parent_uuid` refers to its client call; content-free Messages-shaped output carries known usage and USD charges | + +The summary declares expected attempts in `switchyard.accounting.attempts` +and captured attempts in `switchyard.accounting.attempt_records`. Each child +has a one-based `attempt` ordinal. `outcome` (`ok`, `error`, or `cancelled`) +is separate from `switchyard.accounting.usage_status`. Missing token counts +remain absent; reported zero remains zero. Charges use `usage.cost` with +`currency = "USD"` and `source = "provider_reported"`. +`switchyard.accounting.charge_source` distinguishes body charges (`provider`) +from the supported header fallback (`header`); body charges take precedence. + +`switchyard.usage.input_tokens_reported` preserves whether the original +non-cached input count was known. When true, `switchyard.usage.input_tokens` +contains that count, including zero. When false, the projected input count +stays absent even if cache usage is known. Older captures without this flag +retain the aggregate-minus-cache fallback. + +The example's `coverage_complete` checks the received, instrumented completion +records against their summary counts. It cannot detect a wholly absent call +or certify full workflow coverage. Its `total_cost_usd` has that same scope; +decision calls and other producers require a separate accounting view. + +Relay's [model-pricing guide](https://docs.nvidia.com/nemo/relay/configure-plugins/model-pricing) +explains how to supply model rates. The cost-summary example covers instrumented +completion attempts. Separate decision-model calls are not yet exported with +per-call usage. ATIF `final_metrics` still describes Relay's outer interaction; +it does not automatically include this retry total. Provider work after a +disconnect and missing telemetry can leave costs unknown. ## How Requests Flow @@ -305,8 +359,8 @@ can receive these records, but each output format presents them differently. | Output | What the current integration shows | | --- | --- | | ATIF | The request and response seen by the agent. Internal routing is not added as separate steps or included in `final_metrics`. Relay-managed ATIF files may retain the raw records under `extra.observed_events`. | -| OpenTelemetry traces, including OpenInference | Relay's outer LLM span, nested `libsy.client_call` and `libsy.upstream_attempt` duration scopes, plus eligible routing marks. | -| OpenTelemetry metrics | Switchyard request, routing-call, latency, token, and failure measurements. Combine the reported routing usage with Relay's answer cost to estimate the observed routed cost. | +| OpenTelemetry traces, including OpenInference | Relay's outer LLM span, nested `libsy.client_call` summaries and `libsy.upstream_attempt` accounting scopes, plus eligible routing marks. | +| OpenTelemetry metrics | Switchyard request, routing-call, latency, token, and failure measurements. Aggregate model-call measurements; use attempt scopes for retry cost accounting. | | OTLP logs | Non-metric Switchyard marks that meet the configured severity threshold. | The plugin captures two existing Switchyard spans. Each routing or answer LLM diff --git a/docs/internal/req3_attempt_accounting_implementation.md b/docs/internal/req3_attempt_accounting_implementation.md new file mode 100644 index 0000000000..bcbbd4a996 --- /dev/null +++ b/docs/internal/req3_attempt_accounting_implementation.md @@ -0,0 +1,84 @@ +# Per-attempt accounting implementation + +Implemented on `bbednarski/req3-attempt-accounting`, based on upstream main +`cf92c1135`. The changes are pending review and commit. + +The client retains accounting until each dispatched completion attempt ends. +Its guard preserves usage across HTTP failures, decoding, routing buffering, +streaming errors, timeouts, and cancellation. Usage is cumulative by field. +Missing values remain unknown. Local failures before dispatch have no attempt. + +The native plugin exports priceable attempt scopes, with content-free usage +and bounded error categories. Logical client calls become summary scopes with +expected and captured attempt counts. Existing payload capture remains optional +and supplements the accounting data. + +## Scope + +```mermaid +flowchart TB + Core["libsy / runner algorithms"] --> Client + Protocol["Existing protocol Usage and Response types"] -.-> Guard + Transport["Existing HTTP transport"] <--> Client + subgraph LlmClient["switchyard-llm-client"] + Client["Augment client.rs / run.rs / observability.rs"] --> Guard["NEW attempt.rs guard and stream ownership"] + end + Client <--> Translation["Augment switchyard-translation
Preserve supported failure usage"] + Guard --> Capture["Augment native plugin span_capture.rs
Usage, outcome, summary counts, shared projections"] + Capture --> Relay["Existing Relay SDK and pricing codecs"] + Relay --> Consumer["NEW ATOF consumer example
Unique attempts, priced subtotal, nullable total"] + Consumer --> Fixture["NEW native acceptance fixture"] + classDef unchanged fill:#dcfce7,stroke:#15803d,color:#14532d; + classDef augmented fill:#fef3c7,stroke:#b45309,color:#78350f; + classDef added fill:#fee2e2,stroke:#b91c1c,color:#7f1d1d; + class Core,Protocol,Transport,Relay unchanged; + class Client,Translation,Capture augmented; + class Guard,Consumer,Fixture added; +``` + +Three production crates change. The translation changes reuse provider usage +parsers and retain usage in Responses failed/incomplete terminal frames. +Protocol types, retry policy, and decision transports stay unchanged. + +## Relay integration result + +Tested a packaged macOS arm64 native bundle against `nemo-relay==0.10.0` and +native SDK `0.10.0`, with a loopback provider and fixed synthetic rates. + +Generic native scopes retain their raw output, but do not acquire a normalized +pricing annotation automatically. Relay's outer answer also retains the final +response charge. The smallest follow-up is an opt-in ATOF consumer example that +selects unique attempt end events, decodes their output with Relay's existing +codec, and checks declared record coverage. It adds no price catalog or pricing +formula. No Relay source changes are required for this output. + +Relay 0.10 ignores the scalar `usage.provider_cost` field. The plugin uses the +supported detailed `usage.cost` object so reported charges take precedence over +model-price estimates. + +| Fixture | Known subtotal | Complete total | +| --- | --- | --- | +| Failed reported $0.02, successful reported $0.03 | $0.05 | $0.05 | +| Failed usage missing, successful reported $0.03 | $0.03 | Absent | +| Both attempts priced by configured token rates | $0.028 | $0.028 | +| Usage known, model price missing | $0 | Absent | +| Both attempts report $0, tokens missing | $0 | $0 | + +The fixture also verifies duplicate-event delivery, missing attempt records, +reported-charge precedence over headers, and separation between workflows. +The Rust suites cover routing buffer failure, terminal stream errors, trailing +and repeated usage snapshots, timeout, abandonment, pre-dispatch deadlines, +concurrent requests, debug payload compatibility, and credential redaction. + +## Validation and remaining boundaries + +The full workspace Rust tests, formatting, and Clippy pass. The repository +pre-commit hooks pass. Optional non-blocking mypy/pyright commands are unavailable +in this environment; those hooks report that they did not execute. +The new Python scripts also pass a direct Ruff check and native acceptance. +See the plugin README for the reproducible native fixture command. + +These totals cover observed instrumented completion attempts. Separate decision +calls remain Req 4. ATIF `final_metrics` is unchanged; use the supplied ATOF +summary for retry totals. Provider work continuing after disconnection, process +termination, and wholly absent telemetry cannot be reconciled from these records.