diff --git a/crates/switchyard-translation/src/codecs/anthropic/stream.rs b/crates/switchyard-translation/src/codecs/anthropic/stream.rs index 135fb645c..f958f9cf3 100644 --- a/crates/switchyard-translation/src/codecs/anthropic/stream.rs +++ b/crates/switchyard-translation/src/codecs/anthropic/stream.rs @@ -25,6 +25,10 @@ impl StreamCodec for AnthropicMessagesStreamCodec { event.get("type").and_then(Value::as_str) == Some("message_stop") } + fn stops_at_terminal_event(&self) -> bool { + true + } + fn sse_done_is_terminal(&self) -> bool { false } diff --git a/crates/switchyard-translation/src/codecs/responses/stream.rs b/crates/switchyard-translation/src/codecs/responses/stream.rs index 7a450b535..d2c49754d 100644 --- a/crates/switchyard-translation/src/codecs/responses/stream.rs +++ b/crates/switchyard-translation/src/codecs/responses/stream.rs @@ -37,6 +37,10 @@ impl StreamCodec for OpenAiResponsesStreamCodec { ) } + fn stops_at_terminal_event(&self) -> bool { + true + } + fn sse_done_is_terminal(&self) -> bool { true } diff --git a/crates/switchyard-translation/src/codecs/stream.rs b/crates/switchyard-translation/src/codecs/stream.rs index 145679abf..68d99af6d 100644 --- a/crates/switchyard-translation/src/codecs/stream.rs +++ b/crates/switchyard-translation/src/codecs/stream.rs @@ -172,6 +172,13 @@ pub trait StreamCodec: Send + Sync { /// Returns whether a JSON event establishes stream completion. fn is_terminal_event(&self, event: &Value) -> bool; + /// Returns whether decoding stops after yielding a terminal JSON event. + /// + /// The default keeps reading so codecs can emit trailing events, such as usage. + fn stops_at_terminal_event(&self) -> bool { + false + } + /// Returns whether the SSE `[DONE]` marker establishes stream completion. /// /// The marker ends consumption even when this returns `false`; a terminal JSON event diff --git a/crates/switchyard-translation/src/engine.rs b/crates/switchyard-translation/src/engine.rs index f70d06725..fd715d132 100644 --- a/crates/switchyard-translation/src/engine.rs +++ b/crates/switchyard-translation/src/engine.rs @@ -483,8 +483,9 @@ impl TranslationEngine { /// buffered across chunks (a partial frame waits for its boundary); the source /// stream codec is resolved once and reused for every frame. /// - /// The decoder tracks the source format's protocol-specific terminal event. If - /// EOF arrives without that required event, the stream yields a deferred + /// Codecs can opt in to stopping after the decoder yields a terminal JSON event. + /// Other codecs keep reading to include trailing events before `[DONE]` or EOF. + /// If EOF arrives without a required terminal event, the stream yields a deferred /// [`LlmClientError::ResponseTranslation`] error after any valid decoded events. pub fn decode_stream( &self, @@ -545,6 +546,9 @@ impl TranslationEngine { value, normalized, ); + if saw_terminal && codec.stops_at_terminal_event() { + break; + } } } } else { diff --git a/crates/switchyard-translation/tests/extension_points.rs b/crates/switchyard-translation/tests/extension_points.rs index e3a9f497d..abd3a1856 100644 --- a/crates/switchyard-translation/tests/extension_points.rs +++ b/crates/switchyard-translation/tests/extension_points.rs @@ -155,6 +155,40 @@ fn custom_whole_streams_use_registered_codecs_and_outlive_the_engine() -> TestRe Ok(()) } +// A custom codec can report completion before it sends its last event. +#[test] +fn custom_stream_keeps_events_after_terminal_json() -> TestResult { + let mut registry = StreamCodecRegistry::new(); + registry.register(CustomStreamCodec); + let engine = TranslationEngine::with_registries(FormatRegistry::new(), registry); + let raw = [ + json!({"kind": "done"}), + json!({"kind": "delta", "text": "trailing"}), + ]; + let bytes = stream::iter( + raw.clone() + .map(|event| Ok(format!("data: {event}\n\n").into_bytes())), + ); + let events = block_on( + engine + .decode_stream("custom_stream", bytes)? + .try_collect::>(), + )?; + + assert_eq!(events.len(), 2); + for (event, expected) in events.iter().zip(raw) { + assert_eq!(event.preservation().unwrap().raw(), &expected); + } + assert_eq!( + events[1].normalized(), + &[LlmResponseChunk::TextDelta { + index: 0, + text: "trailing".into(), + }] + ); + Ok(()) +} + #[test] fn custom_stream_completion_uses_the_codec_rules() -> TestResult { let mut registry = StreamCodecRegistry::new(); diff --git a/tests/test_translation_terminal.py b/tests/test_translation_terminal.py new file mode 100644 index 000000000..a52591041 --- /dev/null +++ b/tests/test_translation_terminal.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +import pytest + +from switchyard.translation import TranslationEngine + + +@pytest.mark.parametrize( + ("protocol", "terminal"), + [ + ("anthropic_messages", {"type": "message_stop"}), + ("openai_responses", {"type": "response.completed", "response": {}}), + ("openai_responses", {"type": "response.incomplete", "response": {}}), + ("openai_responses", {"type": "response.failed", "response": {}}), + ], +) +async def test_decoder_stops_at_terminal_event(protocol, terminal): + """The decoder yields the terminal event without polling upstream again.""" + + async def upstream(): + """The iterator raises if the decoder reads beyond the terminal event.""" + yield f"data: {json.dumps(terminal)}\n\n".encode() + raise AssertionError("The decoder read after the terminal event.") + + async with TranslationEngine().decode_stream(protocol, upstream()) as decoded: + assert [event["preservation"]["raw"] async for event in decoded] == [terminal] + + +async def test_chat_decoder_yields_usage_after_finish_reason(): + """The decoder keeps Chat Completions usage that follows finish_reason.""" + usage = {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5} + + async def upstream(): + """The iterator sends usage after finish_reason and raises after [DONE].""" + yield b'data: {"choices":[{"delta":{},"finish_reason":"stop"}]}\n\n' + yield f"data: {json.dumps({'choices': [], 'usage': usage})}\n\n".encode() + yield b"data: [DONE]\n\n" + raise AssertionError("The decoder read after [DONE].") + + async with TranslationEngine().decode_stream("openai_chat", upstream()) as decoded: + events = [event async for event in decoded] + assert events[-1]["preservation"]["raw"]["usage"] == usage + decoded_usage = next(chunk["Usage"] for chunk in events[-1]["normalized"] if "Usage" in chunk) + assert decoded_usage["input_tokens"] == 2 + assert decoded_usage["output_tokens"] == 3 + assert decoded_usage["total_tokens"] == 5