From b0948295c8d4254100abd3216e5bd8c2f8bb2004 Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Sat, 19 Sep 2026 17:07:43 -0700 Subject: [PATCH 1/2] fix(translation): honor Chat streaming usage opt-in Signed-off-by: Anthony Casagrande --- crates/switchyard-server/src/lib.rs | 3 +- .../src/codecs/openai_chat/stream.rs | 50 +++- .../src/codecs/stream.rs | 6 + crates/switchyard-translation/src/helpers.rs | 253 ++++++++++++++---- .../tests/response_translation.rs | 13 +- .../tests/stream_translation.rs | 30 +-- 6 files changed, 277 insertions(+), 78 deletions(-) diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index 58ebc0d97..ad7ff285b 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -1050,7 +1050,8 @@ async fn handle_llm_request( route.algorithm_name(), ) }); - // Only the Codex namespace mapping is needed downstream, not the whole request. + // Response encoding needs caller-owned extensions such as Chat stream usage + // opt-in and Codex tool identity restoration. let request_extensions = request.llm_request.extensions.clone(); let observer = stats_observer( state.stats.clone(), diff --git a/crates/switchyard-translation/src/codecs/openai_chat/stream.rs b/crates/switchyard-translation/src/codecs/openai_chat/stream.rs index 7a902627b..117d5d1a8 100644 --- a/crates/switchyard-translation/src/codecs/openai_chat/stream.rs +++ b/crates/switchyard-translation/src/codecs/openai_chat/stream.rs @@ -38,6 +38,24 @@ impl StreamCodec for OpenAiChatStreamCodec { encode_openai_chat_stream(state, event) } + fn observe_replayed_event( + &self, + state: &mut StreamTranslationState, + _raw: &Value, + normalized: Vec, + ) { + let replayed_terminal = normalized + .iter() + .any(|chunk| matches!(chunk, LlmResponseChunk::MessageStop { .. })); + for chunk in normalized { + drop(self.encode_event(state, chunk)); + } + if replayed_terminal { + state.finished = true; + state.openai_chat_usage_finalized = true; + } + } + fn finish(&self, state: &mut StreamTranslationState) -> Vec { finish_openai_chat_stream(state) } @@ -292,11 +310,7 @@ fn encode_openai_chat_stream( LlmResponseChunk::Usage(usage) => { state.usage = usage; state.saw_backend_usage = true; - if state.finished { - vec![openai_usage_chunk(state)] - } else { - Vec::new() - } + Vec::new() } LlmResponseChunk::MessageStop { reason } => { if state.finished { @@ -307,7 +321,7 @@ fn encode_openai_chat_stream( state, json!({}), Some(openai_finish_reason(reason.as_deref())), - state.saw_backend_usage.then(|| openai_usage_value(state)), + None, )] } LlmResponseChunk::DecodeError { message } | LlmResponseChunk::StreamError { message } => { @@ -321,16 +335,32 @@ fn encode_openai_chat_stream( // Emits a terminal chunk if the source stream ended before a stop event arrived. fn finish_openai_chat_stream(state: &mut StreamTranslationState) -> Vec { - if state.finished || !state.saw_message_start { + if state.errored { + return Vec::new(); + } + if state.finished { + return finalize_openai_chat_usage(state); + } + if !state.saw_message_start { return Vec::new(); } state.finished = true; - vec![openai_stream_chunk( + let mut out = vec![openai_stream_chunk( state, json!({}), Some(openai_finish_reason(state.stop_reason.as_deref())), - state.saw_backend_usage.then(|| openai_usage_value(state)), - )] + None, + )]; + out.extend(finalize_openai_chat_usage(state)); + out +} + +fn finalize_openai_chat_usage(state: &mut StreamTranslationState) -> Vec { + if !state.openai_chat_include_usage || state.openai_chat_usage_finalized { + return Vec::new(); + } + state.openai_chat_usage_finalized = true; + vec![openai_usage_chunk(state)] } // Normalizes OpenAI token usage fields. diff --git a/crates/switchyard-translation/src/codecs/stream.rs b/crates/switchyard-translation/src/codecs/stream.rs index 87da2d649..d607a9f9f 100644 --- a/crates/switchyard-translation/src/codecs/stream.rs +++ b/crates/switchyard-translation/src/codecs/stream.rs @@ -38,6 +38,12 @@ pub struct StreamTranslationState { /// Set once an in-band error event was emitted; the encoder then emits nothing further. pub errored: bool, pub usage: Usage, + /// Whether the inbound Chat caller requested the final usage-only chunk. + #[serde(default)] + pub(crate) openai_chat_include_usage: bool, + /// Marks generated or preserved Chat usage finalization as complete. + #[serde(default)] + pub(crate) openai_chat_usage_finalized: bool, pub(crate) output_tokens_seen: u64, pub(crate) saw_backend_usage: bool, diff --git a/crates/switchyard-translation/src/helpers.rs b/crates/switchyard-translation/src/helpers.rs index 95ebb75e7..c1e68358d 100644 --- a/crates/switchyard-translation/src/helpers.rs +++ b/crates/switchyard-translation/src/helpers.rs @@ -114,10 +114,8 @@ pub fn encode_stream( ) } -/// Encodes a response stream, honouring request extensions. -/// -/// Identical to [`encode_stream`] except that Codex tool namespaces recorded on -/// the request are restored on each encoded event. +/// Encodes a response stream while applying response-shaping state retained +/// from the request, including Chat usage opt-in and Codex tool identities. pub fn encode_stream_with_extensions( chunks: LlmResponseStream, target: WireFormat, @@ -141,6 +139,7 @@ pub fn encode_stream_with_extensions( let mut state = StreamTranslationState { target: Some(target_format.clone()), target_model: served_model, + openai_chat_include_usage: openai_chat_stream_usage_requested(target, request_extensions), ..Default::default() }; let mut chunks = chunks; @@ -200,6 +199,20 @@ pub fn encode_stream_with_extensions( Ok(Box::pin(events)) } +fn openai_chat_stream_usage_requested( + target: WireFormat, + request_extensions: &switchyard_protocol::ProviderExtensions, +) -> bool { + target == WireFormat::OpenAiChat + && request_extensions + .fields + .get("stream_options") + .and_then(Value::as_object) + .and_then(|options| options.get("include_usage")) + .and_then(Value::as_bool) + .unwrap_or(false) +} + // The raw-response helper promises that the caller sees the model that served the // request. Same-format preservation bypasses provider codecs, so apply that // helper-specific override after replay without disturbing any other raw fields. @@ -363,18 +376,67 @@ mod tests { use futures::{Stream, StreamExt, stream}; use serde_json::{Value, json}; use switchyard_protocol::{ - LlmClientError, LlmResponseChunk, LlmResponseStreamEvent, completion_text, + LlmClientError, LlmResponseChunk, LlmResponseStreamEvent, ProviderExtensions, + completion_text, }; use super::{ decode_aggregated_response, decode_request, decode_stream, encode_aggregated_response, - encode_request, encode_stream, stamp_streamed_response_model, + encode_request, encode_stream, encode_stream_with_extensions, + stamp_streamed_response_model, }; use crate::{LlmResponseStream, LlmStreamError, WireFormat}; // A boxed stream item error, matching the streamed IR contract. type BoxError = Box; + fn chat_request_extensions( + include_usage: Option, + ) -> Result { + let mut body = json!({ + "model": "route/model", + "messages": [{"role": "user", "content": "hi"}], + "stream": true, + }); + if let Some(include_usage) = include_usage { + body["stream_options"] = json!({"include_usage": include_usage}); + } + Ok(decode_request(WireFormat::OpenAiChat, &body)?.extensions) + } + + fn collect_chat_events( + chunks: Vec, + extensions: &ProviderExtensions, + ) -> Result, BoxError> { + let chunks: LlmResponseStream = stream::iter( + chunks + .into_iter() + .map(|chunk| Ok::(chunk.into())), + ) + .boxed(); + Ok(block_on( + encode_stream_with_extensions(chunks, WireFormat::OpenAiChat, None, extensions)? + .collect::>(), + ) + .into_iter() + .collect::, LlmStreamError>>()?) + } + + fn usage_chunk() -> LlmResponseChunk { + LlmResponseChunk::Usage(switchyard_protocol::llm::Usage { + input_tokens: Some(10), + output_tokens: Some(5), + total_tokens: Some(15), + ..Default::default() + }) + } + + fn stop_chunk() -> LlmResponseChunk { + LlmResponseChunk::MessageStop { + reason: Some("stop".to_string()), + } + } + // Collects a decoded IR stream, surfacing the first error instead of panicking. fn decode_all( bytes: impl Stream, LlmClientError>> + Send + 'static, @@ -666,45 +728,105 @@ mod tests { Ok(()) } - // The guard keys on `errored`, not `finished`, so a normal completion still emits the - // trailing usage chunk the OpenAI chat codec reports only after `finished` is set. #[test] - fn encode_stream_keeps_trailing_usage_after_a_normal_stop() -> Result<(), BoxError> { - let chunks: LlmResponseStream = stream::iter(vec![ - Ok(LlmResponseChunk::TextDelta { - index: 0, - text: "hi".to_string(), - } - .into()), - Ok(LlmResponseChunk::MessageStop { - reason: Some("stop".to_string()), + fn openai_chat_stream_usage_is_not_generated_without_boolean_opt_in() -> Result<(), BoxError> { + for include_usage in [None, Some(json!(false)), Some(json!("true"))] { + for usage_before_stop in [true, false] { + let extensions = chat_request_extensions(include_usage.clone())?; + let chunks = if usage_before_stop { + vec![usage_chunk(), stop_chunk()] + } else { + vec![stop_chunk(), usage_chunk()] + }; + let events = collect_chat_events(chunks, &extensions)?; + assert!(events.iter().all(|event| event.get("usage").is_none())); + assert_eq!( + events.last().unwrap()["choices"][0]["finish_reason"], + "stop" + ); } - .into()), - Ok(LlmResponseChunk::Usage(switchyard_protocol::llm::Usage { - output_tokens: Some(7), - ..Default::default() + } + Ok(()) + } + + #[test] + fn openai_chat_stream_usage_requested_before_stop_is_trailing() -> Result<(), BoxError> { + let extensions = chat_request_extensions(Some(json!(true)))?; + let events = collect_chat_events(vec![usage_chunk(), stop_chunk()], &extensions)?; + assert_eq!(events.len(), 2); + assert_eq!(events[0]["choices"][0]["finish_reason"], "stop"); + assert!(events[0].get("usage").is_none()); + assert_eq!(events[1]["choices"], json!([])); + assert_eq!( + events[1]["usage"], + json!({ + "prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15 }) - .into()), - ]) - .boxed(); - let events = - block_on(encode_stream(chunks, WireFormat::OpenAiChat, None)?.collect::>()) - .into_iter() - .collect::, LlmStreamError>>()?; - let body = serde_json::to_string(&events)?; - assert!( - events - .iter() - .any(|event| event["choices"][0]["finish_reason"] == "stop"), - "missing stop terminal:\n{body}" ); - assert!( - body.contains("\"usage\""), - "trailing usage dropped after a normal stop:\n{body}" + Ok(()) + } + + #[test] + fn openai_chat_stream_usage_requested_after_stop_is_finalized_at_clean_eof() + -> Result<(), BoxError> { + let extensions = chat_request_extensions(Some(json!(true)))?; + let events = collect_chat_events(vec![stop_chunk(), usage_chunk()], &extensions)?; + assert_eq!(events.len(), 2); + assert_eq!(events[0]["choices"][0]["finish_reason"], "stop"); + assert_eq!(events[1]["choices"], json!([])); + assert_eq!(events[1]["usage"]["total_tokens"], 15); + Ok(()) + } + + #[test] + fn openai_chat_stream_usage_requested_without_backend_usage_uses_zero_defaults() + -> Result<(), BoxError> { + let extensions = chat_request_extensions(Some(json!(true)))?; + let events = collect_chat_events(vec![stop_chunk()], &extensions)?; + assert_eq!(events.len(), 2); + assert_eq!(events[0]["choices"][0]["finish_reason"], "stop"); + assert_eq!(events[1]["choices"], json!([])); + assert_eq!( + events[1]["usage"], + json!({ + "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0 + }) ); Ok(()) } + #[test] + fn openai_chat_stream_usage_is_not_emitted_after_a_post_stop_error() -> Result<(), BoxError> { + let extensions = chat_request_extensions(Some(json!(true)))?; + for usage_before_stop in [true, false] { + let items = if usage_before_stop { + vec![ + Ok::(usage_chunk().into()), + Ok::(stop_chunk().into()), + Err(LlmClientError::General("boom".to_string())), + ] + } else { + vec![ + Ok::(stop_chunk().into()), + Ok::(usage_chunk().into()), + Err(LlmClientError::General("boom".to_string())), + ] + }; + let chunks: LlmResponseStream = stream::iter(items).boxed(); + let results = block_on( + encode_stream_with_extensions(chunks, WireFormat::OpenAiChat, None, &extensions)? + .collect::>(), + ); + + let [Ok(finish), Err(LlmStreamError::Client(_))] = results.as_slice() else { + return Err(format!("expected finish then input error, got {results:?}").into()); + }; + assert_eq!(finish["choices"][0]["finish_reason"], "stop"); + assert!(finish.get("usage").is_none()); + } + Ok(()) + } + #[test] fn decode_stream_parses_sse_bytes_into_ir_chunks() -> Result<(), LlmClientError> { let sse = b"data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\ @@ -719,28 +841,67 @@ mod tests { } #[test] - fn stream_helpers_replay_same_format_provider_fields() -> Result<(), BoxError> { - let provider_event = json!({ + fn openai_chat_stream_usage_does_not_change_same_format_raw_replay() -> Result<(), BoxError> { + let finish = json!({ "id": "chatcmpl-test", "object": "chat.completion.chunk", "system_fingerprint": "fp_provider_specific", "choices": [{ "index": 0, - "delta": {"content": "Hello"}, + "delta": {}, "finish_reason": "stop" }] }); + let usage = json!({ + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_provider_specific", + "choices": [], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} + }); let bytes = stream::once({ - let frame = format!("data: {provider_event}\n\n").into_bytes(); + let frame = format!("data: {finish}\n\ndata: {usage}\n\n").into_bytes(); async move { Ok::, LlmClientError>(frame) } }); let decoded = decode_stream(bytes, WireFormat::OpenAiChat)?; - let replayed = - block_on(encode_stream(decoded, WireFormat::OpenAiChat, None)?.collect::>()) - .into_iter() - .collect::, LlmStreamError>>()?; + let extensions = chat_request_extensions(Some(json!(true)))?; + let replayed = block_on( + encode_stream_with_extensions(decoded, WireFormat::OpenAiChat, None, &extensions)? + .collect::>(), + ) + .into_iter() + .collect::, LlmStreamError>>()?; + + assert_eq!(replayed, vec![finish, usage]); + Ok(()) + } + + #[test] + fn openai_chat_stream_usage_does_not_append_to_finish_only_raw_replay() -> Result<(), BoxError> + { + let finish = json!({ + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": "stop" + }] + }); + let bytes = stream::once({ + let frame = format!("data: {finish}\n\n").into_bytes(); + async move { Ok::, LlmClientError>(frame) } + }); + let decoded = decode_stream(bytes, WireFormat::OpenAiChat)?; + let extensions = chat_request_extensions(Some(json!(true)))?; + let replayed = block_on( + encode_stream_with_extensions(decoded, WireFormat::OpenAiChat, None, &extensions)? + .collect::>(), + ) + .into_iter() + .collect::, LlmStreamError>>()?; - assert_eq!(replayed, vec![provider_event]); + assert_eq!(replayed, vec![finish]); Ok(()) } diff --git a/crates/switchyard-translation/tests/response_translation.rs b/crates/switchyard-translation/tests/response_translation.rs index 4280fa3f1..b000e02d7 100644 --- a/crates/switchyard-translation/tests/response_translation.rs +++ b/crates/switchyard-translation/tests/response_translation.rs @@ -356,18 +356,15 @@ fn responses_reasoning_usage_translates_to_openai_chat_usage_details() -> TestRe WireFormat::OpenAiResponses, WireFormat::OpenAiChat, ); - let events = engine.translate_event( + let mut events = engine.translate_event( &mut state, WireFormat::OpenAiResponses, WireFormat::OpenAiChat, &json!({"type": "response.completed", "response": body}), )?; + events.extend(engine.finish_stream(&mut state, WireFormat::OpenAiChat)?); assert_eq!(state.usage, decoded.usage); - let chat = events - .iter() - .find(|event| event.get("usage").is_some()) - .ok_or("missing usage")?; - assert_eq!(chat["usage"], output["usage"]); + assert!(events.iter().all(|event| event.get("usage").is_none())); let mut state = switchyard_translation::StreamTranslationState::new( WireFormat::OpenAiChat, WireFormat::OpenAiResponses, @@ -615,6 +612,10 @@ fn anthropic_thinking_response_translates_to_openai_reasoning_content() -> TestR assert_eq!(state.usage.reasoning_tokens, Some(thinking_tokens)); assert_eq!(state.usage.output_tokens, Some(7)); events.extend(engine.finish_stream(&mut state, target)?); + if target == WireFormat::OpenAiChat { + assert!(events.iter().all(|event| event.get("usage").is_none())); + continue; + } let usage = events .iter() .find_map(|event| { diff --git a/crates/switchyard-translation/tests/stream_translation.rs b/crates/switchyard-translation/tests/stream_translation.rs index e22511c59..1f286c959 100644 --- a/crates/switchyard-translation/tests/stream_translation.rs +++ b/crates/switchyard-translation/tests/stream_translation.rs @@ -612,8 +612,9 @@ fn anthropic_stream_usage_and_stop_translate_to_openai_chunks() -> TestResult { &usage, )?; - assert_eq!(events[0]["usage"]["completion_tokens"], 42); + assert!(events[0].get("usage").is_none()); assert_eq!(events[0]["choices"][0]["finish_reason"], "stop"); + assert_eq!(state.usage.output_tokens, Some(42)); Ok(()) } @@ -822,9 +823,9 @@ fn openai_chat_finish_synthesizes_terminal_chunk_after_incomplete_source() -> Te Ok(()) } -// Verifies provider usage arriving after finish remains visible to OpenAI clients. +// Verifies provider usage arriving after finish is recorded without an opt-in. #[test] -fn openai_chat_emits_usage_arriving_after_stop() -> TestResult { +fn openai_chat_records_usage_arriving_after_stop_without_opt_in() -> TestResult { let engine = TranslationEngine::default(); let mut state = StreamTranslationState::new(WireFormat::OpenAiChat, WireFormat::OpenAiChat); let stop = json!({ @@ -854,9 +855,8 @@ fn openai_chat_emits_usage_arriving_after_stop() -> TestResult { WireFormat::OpenAiChat, &usage, )?; - assert_eq!(events.len(), 1); - assert_eq!(events[0]["choices"], json!([])); - assert_eq!(events[0]["usage"]["total_tokens"], 15); + assert!(events.is_empty()); + assert_eq!(state.usage.total_tokens, Some(15)); Ok(()) } @@ -1021,15 +1021,15 @@ fn responses_stream_usage_without_total_keeps_cached_tokens_in_total() -> TestRe )?; events.extend(engine.finish_stream(&mut state, WireFormat::OpenAiChat)?); - let Some(usage) = events - .iter() - .find_map(|event| event.get("usage").filter(|usage| !usage.is_null())) - else { - return Err("expected a terminal OpenAI chunk carrying usage".into()); - }; - assert_eq!(usage["prompt_tokens"], 100); - assert_eq!(usage["total_tokens"], 105); - assert_eq!(usage["prompt_tokens_details"]["cached_tokens"], 80); + assert!(events.iter().all(|event| event.get("usage").is_none())); + assert_eq!( + state.usage.input_tokens.unwrap_or(0) + + state.usage.cached_input_tokens().unwrap_or(0) + + state.usage.cache_creation_input_tokens().unwrap_or(0), + 100 + ); + assert_eq!(state.usage.total_tokens, Some(105)); + assert_eq!(state.usage.cached_input_tokens(), Some(80)); Ok(()) } From 71ae90a6b1f710d7abbeaa0ab38d6288c8e20dc2 Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Sat, 19 Sep 2026 17:51:47 -0700 Subject: [PATCH 2/2] test(relay): enforce hidden usage output Signed-off-by: Anthony Casagrande --- crates/switchyard-nemo-relay-plugin/src/runtime.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/switchyard-nemo-relay-plugin/src/runtime.rs b/crates/switchyard-nemo-relay-plugin/src/runtime.rs index 95794f0d8..b43d00e2c 100644 --- a/crates/switchyard-nemo-relay-plugin/src/runtime.rs +++ b/crates/switchyard-nemo-relay-plugin/src/runtime.rs @@ -1538,7 +1538,7 @@ mod tests { assert!(captured.lock().unwrap().is_empty()); assert!(stream.next().await.expect("encoded stream event").is_ok()); assert!(captured.lock().unwrap().is_empty()); - assert!(stream.next().await.expect("encoded usage event").is_ok()); + assert!(stream.next().await.is_none()); let events = captured.lock().unwrap(); assert_eq!(events.len(), 3);