From e3f9b65cf12dd7e3ddbdf796c50155fb39fd7481 Mon Sep 17 00:00:00 2001 From: Elyas Mehtabuddin Date: Mon, 9 Nov 2026 10:24:31 -0800 Subject: [PATCH 1/4] fix(translation): stop reading streams after terminal events Signed-off-by: Elyas Mehtabuddin --- crates/switchyard-translation/src/engine.rs | 11 +- .../tests/extension_points.rs | 5 + tests/test_translation_terminal.py | 201 ++++++++++++++++++ 3 files changed, 214 insertions(+), 3 deletions(-) create mode 100644 tests/test_translation_terminal.py diff --git a/crates/switchyard-translation/src/engine.rs b/crates/switchyard-translation/src/engine.rs index f70d06725..9fb525c93 100644 --- a/crates/switchyard-translation/src/engine.rs +++ b/crates/switchyard-translation/src/engine.rs @@ -22,7 +22,7 @@ use crate::codecs::stream::{ }; use crate::diagnostic::TranslationDiagnostic; use crate::error::{Result, TranslationError}; -use crate::format::FormatId; +use crate::format::{FormatId, WireFormat}; use crate::llm::{AggLlmResponse, LlmRequest, ProviderExtensions}; use crate::policy::TranslationPolicy; use crate::sse; @@ -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 + /// The decoder yields the terminal JSON event before it stops reading. Chat Completions + /// keeps reading after `finish_reason` to include usage 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,10 @@ impl TranslationEngine { value, normalized, ); + // Chat usage can follow finish_reason, so the decoder keeps reading. + if saw_terminal && source_format != WireFormat::OpenAiChat.into() { + break; + } } } } else { diff --git a/crates/switchyard-translation/tests/extension_points.rs b/crates/switchyard-translation/tests/extension_points.rs index e3a9f497d..c2c8cdf3d 100644 --- a/crates/switchyard-translation/tests/extension_points.rs +++ b/crates/switchyard-translation/tests/extension_points.rs @@ -166,6 +166,11 @@ fn custom_stream_completion_uses_the_codec_rules() -> TestResult { ("SSE marker alone", "data: [DONE]\n\n", false), ("trailing SSE marker alone", "data: [DONE]", false), ("terminal JSON event", "data: {\"kind\":\"done\"}\n\n", true), + ( + "terminal JSON event then invalid data", + "data: {\"kind\":\"done\"}\n\ndata: invalid-json\n\n", + true, + ), ( "trailing terminal JSON event", "data: {\"kind\":\"done\"}", diff --git a/tests/test_translation_terminal.py b/tests/test_translation_terminal.py new file mode 100644 index 000000000..dd1b2806f --- /dev/null +++ b/tests/test_translation_terminal.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The decoder stops at protocol completion without waiting for EOF.""" + +import asyncio +import json + +import pytest + +from switchyard.translation import TranslationEngine + +FORMATS = ["openai_chat", "anthropic_messages", "openai_responses"] + + +def source_events(source): + """Chat usage follows finish_reason; Responses text comes from the final snapshot.""" + if source == "openai_chat": + return [ + { + "id": "chat-1", + "model": "backend", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hello"}}], + }, + {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + { + "choices": [], + "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, + }, + ] + if source == "anthropic_messages": + return [ + { + "type": "message_start", + "message": { + "id": "msg-1", + "model": "backend", + "role": "assistant", + "content": [], + "usage": {"input_tokens": 5, "output_tokens": 0}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "hello"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 2}, + }, + {"type": "message_stop"}, + ] + return [ + { + "type": "response.completed", + "response": { + "id": "resp-1", + "model": "backend", + "status": "completed", + "output": [ + { + "id": "msg-1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hello", "annotations": []}], + } + ], + "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, + }, + } + ] + + +def sse(events, source): + payload = b"".join(("data: " + json.dumps(event) + "\n\n").encode() for event in events) + return payload + (b"data: [DONE]\n\n" if source == "openai_chat" else b"") + + +@pytest.mark.parametrize("source", FORMATS) +@pytest.mark.parametrize("target", FORMATS) +@pytest.mark.parametrize("split_tail", [False, True]) +async def test_translation_ignores_data_after_terminal(source, target, split_tail): + payload = sse(source_events(source), source) + tail_event = { + "openai_chat": {"choices": [{"index": 0, "delta": {"content": "AFTER_TERMINAL"}}]}, + "anthropic_messages": { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "AFTER_TERMINAL"}, + }, + "openai_responses": { + "type": "response.output_text.delta", + "item_id": "msg-1", + "output_index": 0, + "content_index": 0, + "delta": "AFTER_TERMINAL", + }, + }[source] + # The invalid frame must be ignored even when it is already buffered. + tail = sse([tail_event], source) + b"data: invalid-json\n\n" + reads = 0 + + async def chunks(): + """The second chunk exposes any read after protocol completion.""" + nonlocal reads + reads += 1 + yield payload if split_tail else payload + tail + reads += 1 + yield tail + + engine = TranslationEngine() + async with engine.translate_stream(source, target, chunks()) as translated: + output = [event async for event in translated] + assert reads == 1 + assert "AFTER_TERMINAL" not in json.dumps(output) + if target == "openai_chat": + assert any( + event.get("choices") and event["choices"][0]["delta"].get("content") == "hello" + for event in output + ) + usage = next(event["usage"] for event in output if "usage" in event) + assert usage["prompt_tokens"] == 5 + assert usage["completion_tokens"] == 2 + elif target == "anthropic_messages": + assert any(event.get("delta", {}).get("text") == "hello" for event in output) + assert output[-1]["type"] == "message_stop" + assert ( + next( + event["usage"]["output_tokens"] + for event in output + if event["type"] == "message_delta" + ) + == 2 + ) + else: + assert output[-1]["type"] == "response.completed" + response = output[-1]["response"] + assert response["output"][0]["content"][0]["text"] == "hello" + assert response["usage"]["input_tokens"] == 5 + assert response["usage"]["output_tokens"] == 2 + + +@pytest.mark.parametrize( + "source,terminal", + [ + ("openai_chat", None), + ("anthropic_messages", {"type": "message_stop"}), + *[ + ( + "openai_responses", + { + key: kind, + "response": { + "id": "resp-1", + "status": kind.removeprefix("response."), + "output": [], + "usage": {"input_tokens": 5, "output_tokens": 2}, + **( + {"error": {"code": "server_error", "message": "provider failed"}} + if kind == "response.failed" + else {} + ), + }, + }, + ) + for key in ("type", "event") + for kind in ("response.completed", "response.incomplete", "response.failed") + ], + ], +) +async def test_decoder_completes_without_reading_after_terminal(source, terminal): + """The terminal payload survives decoding even when the input never reaches EOF.""" + events = source_events(source) if terminal is None else [terminal] + read_after_terminal = False + + async def held_open(): + nonlocal read_after_terminal + # Byte-sized chunks exercise partial frames without reading beyond completion. + for byte in sse(events, source): + yield bytes([byte]) + read_after_terminal = True + await asyncio.Future() + + async def collect(): + async with TranslationEngine().decode_stream(source, held_open()) as decoded: + return [event async for event in decoded] + + output = await asyncio.wait_for(collect(), 2) + assert not read_after_terminal + assert [event["preservation"]["raw"] for event in output] == events + if terminal is not None and terminal.get("type", terminal.get("event")) == "response.failed": + assert {"StreamError": {"message": "provider failed"}} in output[-1]["normalized"] From b83a8c797901a7554798cd85eed661b1a06e8db2 Mon Sep 17 00:00:00 2001 From: Elyas Mehtabuddin Date: Mon, 9 Nov 2026 10:48:47 -0800 Subject: [PATCH 2/4] test(translation): verify terminal stream framing Signed-off-by: Elyas Mehtabuddin --- tests/test_translation_terminal.py | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/tests/test_translation_terminal.py b/tests/test_translation_terminal.py index dd1b2806f..413c4e5b6 100644 --- a/tests/test_translation_terminal.py +++ b/tests/test_translation_terminal.py @@ -14,7 +14,11 @@ def source_events(source): - """Chat usage follows finish_reason; Responses text comes from the final snapshot.""" + """Chat usage follows finish_reason; Responses text comes from the final snapshot. + + Responses starts with response.created so translation can emit its start event + before the final snapshot supplies text. + """ if source == "openai_chat": return [ { @@ -59,6 +63,15 @@ def source_events(source): {"type": "message_stop"}, ] return [ + { + "type": "response.created", + "response": { + "id": "resp-1", + "model": "backend", + "status": "in_progress", + "output": [], + }, + }, { "type": "response.completed", "response": { @@ -76,7 +89,7 @@ def source_events(source): ], "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, }, - } + }, ] @@ -131,17 +144,14 @@ async def chunks(): assert usage["prompt_tokens"] == 5 assert usage["completion_tokens"] == 2 elif target == "anthropic_messages": + assert output[0]["type"] == "message_start" assert any(event.get("delta", {}).get("text") == "hello" for event in output) assert output[-1]["type"] == "message_stop" - assert ( - next( - event["usage"]["output_tokens"] - for event in output - if event["type"] == "message_delta" - ) - == 2 - ) + usage = next(event["usage"] for event in output if event["type"] == "message_delta") + assert usage["output_tokens"] == 2 + assert usage.get("input_tokens", output[0]["message"]["usage"]["input_tokens"]) == 5 else: + assert output[0]["type"] == "response.created" assert output[-1]["type"] == "response.completed" response = output[-1]["response"] assert response["output"][0]["content"][0]["text"] == "hello" From c6f969f99bbe35bb8befe93865fa9fbd503bf1ad Mon Sep 17 00:00:00 2001 From: Elyas Mehtabuddin Date: Mon, 9 Nov 2026 11:45:05 -0800 Subject: [PATCH 3/4] test(translation): keep one concise stream regression Signed-off-by: Elyas Mehtabuddin --- .../tests/extension_points.rs | 5 - tests/test_translation_terminal.py | 210 +----------------- 2 files changed, 7 insertions(+), 208 deletions(-) diff --git a/crates/switchyard-translation/tests/extension_points.rs b/crates/switchyard-translation/tests/extension_points.rs index c2c8cdf3d..e3a9f497d 100644 --- a/crates/switchyard-translation/tests/extension_points.rs +++ b/crates/switchyard-translation/tests/extension_points.rs @@ -166,11 +166,6 @@ fn custom_stream_completion_uses_the_codec_rules() -> TestResult { ("SSE marker alone", "data: [DONE]\n\n", false), ("trailing SSE marker alone", "data: [DONE]", false), ("terminal JSON event", "data: {\"kind\":\"done\"}\n\n", true), - ( - "terminal JSON event then invalid data", - "data: {\"kind\":\"done\"}\n\ndata: invalid-json\n\n", - true, - ), ( "trailing terminal JSON event", "data: {\"kind\":\"done\"}", diff --git a/tests/test_translation_terminal.py b/tests/test_translation_terminal.py index 413c4e5b6..717588761 100644 --- a/tests/test_translation_terminal.py +++ b/tests/test_translation_terminal.py @@ -1,211 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The decoder stops at protocol completion without waiting for EOF.""" - -import asyncio -import json - -import pytest - from switchyard.translation import TranslationEngine -FORMATS = ["openai_chat", "anthropic_messages", "openai_responses"] - -def source_events(source): - """Chat usage follows finish_reason; Responses text comes from the final snapshot. +async def test_decoder_stops_at_message_stop(): + async def upstream(): + yield b'data: {"type":"message_stop"}\n\n' + raise AssertionError("decoder read after message_stop") - Responses starts with response.created so translation can emit its start event - before the final snapshot supplies text. - """ - if source == "openai_chat": - return [ - { - "id": "chat-1", - "model": "backend", - "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hello"}}], - }, - {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, - { - "choices": [], - "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, - }, + async with TranslationEngine().decode_stream("anthropic_messages", upstream()) as decoded: + assert [event["preservation"]["raw"] async for event in decoded] == [ + {"type": "message_stop"} ] - if source == "anthropic_messages": - return [ - { - "type": "message_start", - "message": { - "id": "msg-1", - "model": "backend", - "role": "assistant", - "content": [], - "usage": {"input_tokens": 5, "output_tokens": 0}, - }, - }, - { - "type": "content_block_start", - "index": 0, - "content_block": {"type": "text", "text": ""}, - }, - { - "type": "content_block_delta", - "index": 0, - "delta": {"type": "text_delta", "text": "hello"}, - }, - {"type": "content_block_stop", "index": 0}, - { - "type": "message_delta", - "delta": {"stop_reason": "end_turn"}, - "usage": {"output_tokens": 2}, - }, - {"type": "message_stop"}, - ] - return [ - { - "type": "response.created", - "response": { - "id": "resp-1", - "model": "backend", - "status": "in_progress", - "output": [], - }, - }, - { - "type": "response.completed", - "response": { - "id": "resp-1", - "model": "backend", - "status": "completed", - "output": [ - { - "id": "msg-1", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "hello", "annotations": []}], - } - ], - "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, - }, - }, - ] - - -def sse(events, source): - payload = b"".join(("data: " + json.dumps(event) + "\n\n").encode() for event in events) - return payload + (b"data: [DONE]\n\n" if source == "openai_chat" else b"") - - -@pytest.mark.parametrize("source", FORMATS) -@pytest.mark.parametrize("target", FORMATS) -@pytest.mark.parametrize("split_tail", [False, True]) -async def test_translation_ignores_data_after_terminal(source, target, split_tail): - payload = sse(source_events(source), source) - tail_event = { - "openai_chat": {"choices": [{"index": 0, "delta": {"content": "AFTER_TERMINAL"}}]}, - "anthropic_messages": { - "type": "content_block_delta", - "index": 0, - "delta": {"type": "text_delta", "text": "AFTER_TERMINAL"}, - }, - "openai_responses": { - "type": "response.output_text.delta", - "item_id": "msg-1", - "output_index": 0, - "content_index": 0, - "delta": "AFTER_TERMINAL", - }, - }[source] - # The invalid frame must be ignored even when it is already buffered. - tail = sse([tail_event], source) + b"data: invalid-json\n\n" - reads = 0 - - async def chunks(): - """The second chunk exposes any read after protocol completion.""" - nonlocal reads - reads += 1 - yield payload if split_tail else payload + tail - reads += 1 - yield tail - - engine = TranslationEngine() - async with engine.translate_stream(source, target, chunks()) as translated: - output = [event async for event in translated] - assert reads == 1 - assert "AFTER_TERMINAL" not in json.dumps(output) - if target == "openai_chat": - assert any( - event.get("choices") and event["choices"][0]["delta"].get("content") == "hello" - for event in output - ) - usage = next(event["usage"] for event in output if "usage" in event) - assert usage["prompt_tokens"] == 5 - assert usage["completion_tokens"] == 2 - elif target == "anthropic_messages": - assert output[0]["type"] == "message_start" - assert any(event.get("delta", {}).get("text") == "hello" for event in output) - assert output[-1]["type"] == "message_stop" - usage = next(event["usage"] for event in output if event["type"] == "message_delta") - assert usage["output_tokens"] == 2 - assert usage.get("input_tokens", output[0]["message"]["usage"]["input_tokens"]) == 5 - else: - assert output[0]["type"] == "response.created" - assert output[-1]["type"] == "response.completed" - response = output[-1]["response"] - assert response["output"][0]["content"][0]["text"] == "hello" - assert response["usage"]["input_tokens"] == 5 - assert response["usage"]["output_tokens"] == 2 - - -@pytest.mark.parametrize( - "source,terminal", - [ - ("openai_chat", None), - ("anthropic_messages", {"type": "message_stop"}), - *[ - ( - "openai_responses", - { - key: kind, - "response": { - "id": "resp-1", - "status": kind.removeprefix("response."), - "output": [], - "usage": {"input_tokens": 5, "output_tokens": 2}, - **( - {"error": {"code": "server_error", "message": "provider failed"}} - if kind == "response.failed" - else {} - ), - }, - }, - ) - for key in ("type", "event") - for kind in ("response.completed", "response.incomplete", "response.failed") - ], - ], -) -async def test_decoder_completes_without_reading_after_terminal(source, terminal): - """The terminal payload survives decoding even when the input never reaches EOF.""" - events = source_events(source) if terminal is None else [terminal] - read_after_terminal = False - - async def held_open(): - nonlocal read_after_terminal - # Byte-sized chunks exercise partial frames without reading beyond completion. - for byte in sse(events, source): - yield bytes([byte]) - read_after_terminal = True - await asyncio.Future() - - async def collect(): - async with TranslationEngine().decode_stream(source, held_open()) as decoded: - return [event async for event in decoded] - - output = await asyncio.wait_for(collect(), 2) - assert not read_after_terminal - assert [event["preservation"]["raw"] for event in output] == events - if terminal is not None and terminal.get("type", terminal.get("event")) == "response.failed": - assert {"StreamError": {"message": "provider failed"}} in output[-1]["normalized"] From bc85aad1b9936db30604cfdee9bd469572e9d250 Mon Sep 17 00:00:00 2001 From: Elyas Mehtabuddin Date: Mon, 9 Nov 2026 13:33:47 -0800 Subject: [PATCH 4/4] fix(translation): let codecs opt in to terminal stopping Signed-off-by: Elyas Mehtabuddin --- .../src/codecs/anthropic/stream.rs | 4 ++ .../src/codecs/responses/stream.rs | 4 ++ .../src/codecs/stream.rs | 7 +++ crates/switchyard-translation/src/engine.rs | 9 ++-- .../tests/extension_points.rs | 34 +++++++++++++ tests/test_translation_terminal.py | 48 ++++++++++++++++--- 6 files changed, 94 insertions(+), 12 deletions(-) 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 9fb525c93..fd715d132 100644 --- a/crates/switchyard-translation/src/engine.rs +++ b/crates/switchyard-translation/src/engine.rs @@ -22,7 +22,7 @@ use crate::codecs::stream::{ }; use crate::diagnostic::TranslationDiagnostic; use crate::error::{Result, TranslationError}; -use crate::format::{FormatId, WireFormat}; +use crate::format::FormatId; use crate::llm::{AggLlmResponse, LlmRequest, ProviderExtensions}; use crate::policy::TranslationPolicy; use crate::sse; @@ -483,8 +483,8 @@ 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 yields the terminal JSON event before it stops reading. Chat Completions - /// keeps reading after `finish_reason` to include usage before `[DONE]` or EOF. + /// 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( @@ -546,8 +546,7 @@ impl TranslationEngine { value, normalized, ); - // Chat usage can follow finish_reason, so the decoder keeps reading. - if saw_terminal && source_format != WireFormat::OpenAiChat.into() { + if saw_terminal && codec.stops_at_terminal_event() { break; } } 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 index 717588761..a52591041 100644 --- a/tests/test_translation_terminal.py +++ b/tests/test_translation_terminal.py @@ -1,15 +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 -async def test_decoder_stops_at_message_stop(): +@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(): - yield b'data: {"type":"message_stop"}\n\n' - raise AssertionError("decoder read after message_stop") + """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("anthropic_messages", upstream()) as decoded: - assert [event["preservation"]["raw"] async for event in decoded] == [ - {"type": "message_stop"} - ] + 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