switchyard_translation/
helpers.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Convenience wrappers over the default [`TranslationEngine`] — decode a wire
5//! request/response to the neutral IR, encode the IR back, and decode/encode a
6//! streamed response — so callers can translate without threading an engine and
7//! policy through every call.
8
9use std::pin::Pin;
10use std::sync::LazyLock;
11
12use async_stream::try_stream;
13use futures::io::AsyncBufReadExt;
14use futures::{Stream, StreamExt, TryStreamExt};
15use serde_json::Value;
16use switchyard_protocol::LlmClientError;
17
18use crate::codecs::stream::encode_response_stream_event;
19use crate::sse;
20use crate::{
21    AggLlmResponse, FormatId, LlmRequest, LlmResponseChunk, LlmResponseStream,
22    LlmResponseStreamEvent, LlmStreamError, Result, StreamCodecRegistry, StreamTranslationState,
23    TranslationEngine, TranslationPolicy, WireFormat,
24};
25
26static DEFAULT_TRANSLATION_POLICY: LazyLock<TranslationPolicy> =
27    LazyLock::new(TranslationPolicy::default);
28static DEFAULT_TRANSLATION_ENGINE: LazyLock<TranslationEngine> =
29    LazyLock::new(TranslationEngine::default);
30
31/// Decodes a `wire_format` request body into the neutral IR.
32pub fn decode_request(wire_format: WireFormat, body: &Value) -> Result<LlmRequest> {
33    Ok(DEFAULT_TRANSLATION_ENGINE
34        .decode_request(wire_format, body, &DEFAULT_TRANSLATION_POLICY)?
35        .request)
36}
37
38/// Encodes a normalized request into `wire_format`'s JSON body.
39pub fn encode_request(request: &LlmRequest, wire_format: WireFormat) -> Result<Value> {
40    Ok(DEFAULT_TRANSLATION_ENGINE
41        .encode_request(wire_format, request, &DEFAULT_TRANSLATION_POLICY)?
42        .body)
43}
44
45/// Decodes a buffered `wire_format` response body into the neutral aggregate.
46pub fn decode_aggregated_response(body: &Value, wire_format: WireFormat) -> Result<AggLlmResponse> {
47    Ok(DEFAULT_TRANSLATION_ENGINE
48        .decode_response(wire_format, body, &DEFAULT_TRANSLATION_POLICY)?
49        .response)
50}
51
52/// Encodes a buffered aggregate into `wire_format`'s JSON body, stamping
53/// `served_model` over the encoded id so the caller sees which model answered.
54/// Passing `None` leaves the id the upstream reported.
55pub fn encode_aggregated_response(
56    agg: &AggLlmResponse,
57    wire_format: WireFormat,
58    served_model: Option<&str>,
59) -> Result<Value> {
60    encode_aggregated_response_with_extensions(
61        agg,
62        wire_format,
63        served_model,
64        &switchyard_protocol::ProviderExtensions::default(),
65    )
66}
67
68/// Encodes an aggregated response, honouring request extensions.
69///
70/// Identical to [`encode_aggregated_response`] except that Codex tool
71/// namespaces recorded on the request are restored on the encoded body.
72pub fn encode_aggregated_response_with_extensions(
73    agg: &AggLlmResponse,
74    wire_format: WireFormat,
75    served_model: Option<&str>,
76    request_extensions: &switchyard_protocol::ProviderExtensions,
77) -> Result<Value> {
78    let mut body = DEFAULT_TRANSLATION_ENGINE
79        .encode_response_with_extensions(
80            wire_format,
81            agg,
82            request_extensions,
83            &DEFAULT_TRANSLATION_POLICY,
84        )?
85        .body;
86    if let (Some(model), Value::Object(object)) = (served_model, &mut body) {
87        object.insert("model".to_string(), Value::String(model.to_string()));
88    }
89    Ok(body)
90}
91
92/// A stream of wire-format event objects in one format — the unframed body of an
93/// SSE response. The serving layer frames each `Value` (e.g. as an SSE
94/// `data:`/`event:` block).
95pub type RawEventStream =
96    Pin<Box<dyn Stream<Item = std::result::Result<Value, LlmStreamError>> + Send>>;
97
98/// Encodes a stream of IR chunks into a stream of target-format wire events.
99///
100/// `served_model` is exposed as the response model (via the stream state's
101/// `target_model`); `None` falls back to the id observed on the source stream.
102/// The target stream codec is resolved once and reused per chunk; terminal events
103/// (`message_stop` / `response.completed`) come from `finish`.
104pub fn encode_stream(
105    chunks: LlmResponseStream,
106    target: WireFormat,
107    served_model: Option<String>,
108) -> std::result::Result<RawEventStream, LlmClientError> {
109    encode_stream_with_extensions(
110        chunks,
111        target,
112        served_model,
113        &switchyard_protocol::ProviderExtensions::default(),
114    )
115}
116
117/// Encodes a response stream, honouring request extensions.
118///
119/// Identical to [`encode_stream`] except that Codex tool namespaces recorded on
120/// the request are restored on each encoded event.
121pub fn encode_stream_with_extensions(
122    chunks: LlmResponseStream,
123    target: WireFormat,
124    served_model: Option<String>,
125    request_extensions: &switchyard_protocol::ProviderExtensions,
126) -> std::result::Result<RawEventStream, LlmClientError> {
127    let origins = crate::codex_namespaces::qualified_tool_origins(request_extensions);
128    let custom_tools = crate::codex_custom_tools::custom_tool_names(request_extensions);
129    let mut custom_state = crate::codex_custom_tools::CustomToolCallStreamState::default();
130    let target_format: FormatId = target.into();
131    // The target is always a built-in wire format, so this lookup cannot fail; a
132    // failure returns as an `Err` rather than a panic.
133    let codec = StreamCodecRegistry::with_builtins()
134        .codec(target_format.clone())
135        // Currently the only error is that the codec is missing, which is Configuration
136        .map_err(|err| LlmClientError::Configuration {
137            message: err.to_string(),
138        })?;
139
140    let served_model_for_events = served_model.clone();
141    let mut state = StreamTranslationState {
142        target: Some(target_format.clone()),
143        target_model: served_model,
144        ..Default::default()
145    };
146    let mut chunks = chunks;
147
148    let events = try_stream! {
149        while let Some(item) = chunks.next().await {
150            let event = item.map_err(LlmStreamError::Client)?;
151            let mut encoded =
152                encode_response_stream_event(&mut state, codec.as_ref(), &target_format, event);
153            for value in &mut encoded {
154                stamp_streamed_response_model(value, target, served_model_for_events.as_deref());
155                crate::codex_namespaces::restore_qualified_tool_names(value, &origins);
156            }
157            // Argument deltas for a freeform tool cannot be expressed on the wire; the
158            // rewritten completed item carries the input instead.
159            let mut encoded: Vec<Value> = encoded
160                .into_iter()
161                .filter_map(|mut value| {
162                    crate::codex_custom_tools::restore_custom_tool_calls_in_event(
163                        &mut value,
164                        &custom_tools,
165                        &mut custom_state,
166                    )
167                    .then_some(value)
168                })
169                .collect();
170            let terminal = encoded.pop();
171            for value in encoded {
172                yield value;
173            }
174            match (state.errored, terminal) {
175                (false, Some(value)) => yield value,
176                (true, Some(value)) => Err(LlmStreamError::Upstream(value))?,
177                (true, None) => Err(LlmStreamError::Client(LlmClientError::ResponseTranslation(
178                    "stream ended after an error without a terminal event".to_string(),
179                )))?,
180                _ => ()
181            }
182        }
183        for mut value in codec.finish(&mut state) {
184            stamp_streamed_response_model(
185                &mut value,
186                target,
187                served_model_for_events.as_deref(),
188            );
189            crate::codex_namespaces::restore_qualified_tool_names(&mut value, &origins);
190            if crate::codex_custom_tools::restore_custom_tool_calls_in_event(
191                &mut value,
192                &custom_tools,
193                &mut custom_state,
194            ) {
195                yield value;
196            }
197        }
198    };
199
200    Ok(Box::pin(events))
201}
202
203// The raw-response helper promises that the caller sees the model that served the
204// request. Same-format preservation bypasses provider codecs, so apply that
205// helper-specific override after replay without disturbing any other raw fields.
206fn stamp_streamed_response_model(
207    event: &mut Value,
208    target: WireFormat,
209    served_model: Option<&str>,
210) {
211    let Some(served_model) = served_model else {
212        return;
213    };
214
215    match target {
216        WireFormat::OpenAiChat => {
217            if let Some(event) = event.as_object_mut() {
218                event.insert("model".to_string(), Value::String(served_model.to_string()));
219            }
220        }
221        WireFormat::OpenAiResponses => {
222            if let Some(response) = event.get_mut("response").and_then(Value::as_object_mut) {
223                response.insert("model".to_string(), Value::String(served_model.to_string()));
224            }
225        }
226        WireFormat::AnthropicMessages => {
227            if let Some(message) = event.get_mut("message").and_then(Value::as_object_mut) {
228                message.insert("model".to_string(), Value::String(served_model.to_string()));
229            }
230        }
231    }
232}
233
234/// Decodes provider SSE bytes into normalized stream events.
235///
236/// Operates on raw bytes, not any HTTP client type: the caller adapts its
237/// transport's body stream into `Stream<Item = Result<Vec<u8>, _>>`. Frames are
238/// buffered across chunks (a partial frame waits for its boundary); the source
239/// stream codec is resolved once and reused for every frame.
240///
241/// The decoder tracks the source format's protocol-specific terminal event. If
242/// EOF arrives without that required event, the stream yields a deferred
243/// [`LlmClientError::ResponseTranslation`] error after any valid decoded events.
244pub fn decode_stream<S>(
245    bytes: S,
246    source: WireFormat,
247) -> std::result::Result<LlmResponseStream, LlmClientError>
248where
249    S: Stream<Item = std::result::Result<Vec<u8>, LlmClientError>> + Send + 'static,
250{
251    let marker = sse::done_marker(source);
252    let source_format: FormatId = source.into();
253    // The source is always a built-in wire format, so this lookup cannot fail; a
254    // failure returns as an `Err` rather than a panic.
255    let codec = StreamCodecRegistry::with_builtins()
256        .codec(source_format.clone())
257        .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?;
258    // Adapt the byte-chunk stream into an async line reader. The BufReader
259    // reassembles data split across network chunks (including multi-byte UTF-8),
260    // and `lines()` yields one SSE field line at a time. The stream is boxed to
261    // an `io::Error` item so `into_async_read`'s error bound resolves cleanly. The
262    // source error is boxed intact rather than stringified, so
263    // `llm_client_error_from_io` can recover its original variant on the way out.
264    let io_bytes: Pin<Box<dyn Stream<Item = std::io::Result<Vec<u8>>> + Send>> =
265        Box::pin(bytes.map(|item| item.map_err(std::io::Error::other)));
266    let lines = futures::io::BufReader::new(io_bytes.into_async_read()).lines();
267
268    let mut state = StreamTranslationState {
269        source: Some(source_format.clone()),
270        ..StreamTranslationState::default()
271    };
272    let mut frame = String::new();
273    let mut saw_terminal = false;
274    let mut saw_error = false;
275    let stream = Box::pin(try_stream! {
276        futures::pin_mut!(lines);
277        while let Some(line) = lines.next().await {
278            let line = line.map_err(llm_client_error_from_io)?;
279            // A blank line (allowing a bare CR for CRLF streams) ends the frame.
280            if line.trim_end().is_empty() {
281                let parsed = sse::parse_json_sse_frame(&frame, marker)
282                    .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?;
283                frame.clear();
284                match parsed {
285                    sse::SseFrame::Empty => {}
286                    sse::SseFrame::Done => {
287                        saw_terminal |= source != WireFormat::AnthropicMessages;
288                        break;
289                    }
290                    sse::SseFrame::Data(value) => {
291                        saw_terminal |= sse::is_terminal_event(source, &value);
292                        let normalized = codec.decode_event(&mut state, &value);
293                        saw_error |= normalized.iter().any(|chunk| matches!(
294                            chunk,
295                            LlmResponseChunk::DecodeError { .. }
296                                | LlmResponseChunk::StreamError { .. }
297                        ));
298                        yield LlmResponseStreamEvent::preserved(
299                            source_format.clone(),
300                            value,
301                            normalized,
302                        );
303                    }
304                }
305            } else {
306                frame.push_str(&line);
307                frame.push('\n');
308            }
309        }
310
311        // A non-standard upstream might omit the final blank line; parse a trailing
312        // complete frame instead of losing its last chunk.
313        #[allow(clippy::collapsible_if)]
314        if !frame.trim_end().is_empty() {
315            let parsed = sse::parse_json_sse_frame(&frame, marker)
316                .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?;
317            match parsed {
318                sse::SseFrame::Done => {
319                    saw_terminal |= source != WireFormat::AnthropicMessages;
320                }
321                sse::SseFrame::Data(value) => {
322                    saw_terminal |= sse::is_terminal_event(source, &value);
323                    let normalized = codec.decode_event(&mut state, &value);
324                    saw_error |= normalized.iter().any(|chunk| matches!(
325                        chunk,
326                        LlmResponseChunk::DecodeError { .. }
327                            | LlmResponseChunk::StreamError { .. }
328                    ));
329                    yield LlmResponseStreamEvent::preserved(source_format, value, normalized);
330                }
331                sse::SseFrame::Empty => {}
332            }
333        }
334
335        if !saw_terminal && !saw_error {
336            Err(LlmClientError::ResponseTranslation(format!(
337                "{source} stream ended before a terminal event"
338            )))?;
339        }
340    });
341    Ok(stream)
342}
343
344// Recover transport errors wrapped for `AsyncRead`; other reader failures are
345// invalid upstream responses.
346fn llm_client_error_from_io(error: std::io::Error) -> LlmClientError {
347    let kind = error.kind();
348    let message = error.to_string();
349    match error.into_inner() {
350        Some(source) => match source.downcast::<LlmClientError>() {
351            Ok(error) => *error,
352            Err(source) => LlmClientError::InvalidResponse { source },
353        },
354        None => LlmClientError::InvalidResponse {
355            source: Box::new(std::io::Error::new(kind, message)),
356        },
357    }
358}
359
360#[cfg(test)]
361mod tests {
362    use futures::executor::block_on;
363    use futures::{Stream, StreamExt, stream};
364    use serde_json::{Value, json};
365    use switchyard_protocol::{
366        LlmClientError, LlmResponseChunk, LlmResponseStreamEvent, completion_text,
367    };
368
369    use super::{
370        decode_aggregated_response, decode_request, decode_stream, encode_aggregated_response,
371        encode_request, encode_stream, stamp_streamed_response_model,
372    };
373    use crate::{LlmResponseStream, LlmStreamError, WireFormat};
374
375    // A boxed stream item error, matching the streamed IR contract.
376    type BoxError = Box<dyn std::error::Error + Send + Sync>;
377
378    // Collects a decoded IR stream, surfacing the first error instead of panicking.
379    fn decode_all(
380        bytes: impl Stream<Item = Result<Vec<u8>, LlmClientError>> + Send + 'static,
381        source: WireFormat,
382    ) -> Result<Vec<LlmResponseStreamEvent>, LlmClientError> {
383        block_on(decode_stream(bytes, source)?.collect::<Vec<_>>())
384            .into_iter()
385            .collect()
386    }
387
388    // Concatenates the text of every `TextDelta` chunk.
389    fn text_of(events: &[LlmResponseStreamEvent]) -> String {
390        events
391            .iter()
392            .flat_map(LlmResponseStreamEvent::normalized)
393            .filter_map(|chunk| {
394                if let LlmResponseChunk::TextDelta { text, .. } = chunk {
395                    Some(text.as_str())
396                } else {
397                    None
398                }
399            })
400            .collect()
401    }
402
403    #[test]
404    fn request_round_trips_through_openai_chat() -> Result<(), BoxError> {
405        let body = json!({"model": "gpt", "messages": [{"role": "user", "content": "hi"}]});
406        let request = decode_request(WireFormat::OpenAiChat, &body)?;
407        assert_eq!(request.model.as_deref(), Some("gpt"));
408
409        let encoded = encode_request(&request, WireFormat::OpenAiChat)?;
410        assert_eq!(encoded["model"], "gpt");
411        assert_eq!(encoded["messages"][0]["content"], "hi");
412        Ok(())
413    }
414
415    #[test]
416    fn aggregated_response_round_trips_and_stamps_the_served_model() -> Result<(), BoxError> {
417        let body = json!({
418            "id": "1",
419            "model": "upstream",
420            "choices": [{
421                "index": 0,
422                "message": {"role": "assistant", "content": "Hi there"},
423                "finish_reason": "stop"
424            }],
425            "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}
426        });
427        let agg = decode_aggregated_response(&body, WireFormat::OpenAiChat)?;
428        assert_eq!(completion_text(&agg), "Hi there");
429
430        // `served_model` overrides the id the upstream reported.
431        let encoded =
432            encode_aggregated_response(&agg, WireFormat::OpenAiChat, Some("served/model"))?;
433        assert_eq!(encoded["model"], "served/model");
434        assert_eq!(encoded["choices"][0]["message"]["content"], "Hi there");
435        Ok(())
436    }
437
438    #[test]
439    fn encode_stream_reassembles_deltas_and_finishes() -> Result<(), BoxError> {
440        let chunks: LlmResponseStream = stream::iter(vec![
441            Ok(LlmResponseChunk::TextDelta {
442                index: 0,
443                text: "Hello".to_string(),
444            }
445            .into()),
446            Ok(LlmResponseChunk::TextDelta {
447                index: 0,
448                text: " world".to_string(),
449            }
450            .into()),
451            Ok(LlmResponseChunk::MessageStop {
452                reason: Some("stop".to_string()),
453            }
454            .into()),
455        ])
456        .boxed();
457
458        let events = block_on(
459            encode_stream(chunks, WireFormat::OpenAiChat, Some("m".to_string()))?
460                .collect::<Vec<_>>(),
461        )
462        .into_iter()
463        .collect::<Result<Vec<Value>, LlmStreamError>>()?;
464
465        let content: String = events
466            .iter()
467            .filter_map(|event| event["choices"][0]["delta"]["content"].as_str())
468            .collect();
469        assert_eq!(content, "Hello world");
470        // The terminal chunk carries the stop reason through to the wire events.
471        assert!(
472            events
473                .iter()
474                .any(|event| event["choices"][0]["finish_reason"] == "stop")
475        );
476        Ok(())
477    }
478
479    // The served model must win over the id the source stream announced, so a
480    // routed response never reports the route the caller addressed.
481    #[test]
482    fn encode_stream_stamps_the_served_model_on_message_start() -> Result<(), BoxError> {
483        let chunks: LlmResponseStream = stream::iter(vec![
484            Ok(LlmResponseChunk::MessageStart {
485                id: Some("msg_1".to_string()),
486                model: Some("upstream/model".to_string()),
487            }
488            .into()),
489            Ok(LlmResponseChunk::TextDelta {
490                index: 0,
491                text: "hi".to_string(),
492            }
493            .into()),
494        ])
495        .boxed();
496
497        let events = block_on(
498            encode_stream(
499                chunks,
500                WireFormat::AnthropicMessages,
501                Some("served/model".to_string()),
502            )?
503            .collect::<Vec<_>>(),
504        )
505        .into_iter()
506        .collect::<Result<Vec<Value>, LlmStreamError>>()?;
507
508        assert_eq!(events[0]["type"], "message_start");
509        assert_eq!(events[0]["message"]["model"], "served/model");
510        Ok(())
511    }
512
513    // Without a served model the upstream id survives, so passthrough routes keep
514    // reporting whatever the provider announced.
515    #[test]
516    fn encode_stream_falls_back_to_the_source_model() -> Result<(), BoxError> {
517        let chunks: LlmResponseStream = stream::iter(vec![
518            Ok(LlmResponseChunk::MessageStart {
519                id: Some("msg_1".to_string()),
520                model: Some("upstream/model".to_string()),
521            }
522            .into()),
523            Ok(LlmResponseChunk::TextDelta {
524                index: 0,
525                text: "hi".to_string(),
526            }
527            .into()),
528        ])
529        .boxed();
530
531        let events = block_on(
532            encode_stream(chunks, WireFormat::AnthropicMessages, None)?.collect::<Vec<_>>(),
533        )
534        .into_iter()
535        .collect::<Result<Vec<Value>, LlmStreamError>>()?;
536
537        assert_eq!(events[0]["message"]["model"], "upstream/model");
538        Ok(())
539    }
540
541    #[test]
542    fn encode_stream_propagates_chunk_errors() -> Result<(), BoxError> {
543        let chunks: LlmResponseStream =
544            stream::iter(vec![Err::<LlmResponseStreamEvent, LlmClientError>(
545                LlmClientError::General("chunk exploded".to_string()),
546            )])
547            .boxed();
548        let results =
549            block_on(encode_stream(chunks, WireFormat::OpenAiChat, None)?.collect::<Vec<_>>());
550        assert!(results.iter().any(Result::is_err));
551        Ok(())
552    }
553
554    // An in-band error is terminal for every target format: the encoder emits the pre-error
555    // content, surfaces the error event through the error channel, and drops any later chunk.
556    // Production truncates the source before the encoder, so this contract is only observable
557    // by driving encode_stream directly.
558    #[test]
559    fn encode_stream_stops_after_an_in_band_error() -> Result<(), BoxError> {
560        for message in [
561            LlmResponseChunk::StreamError {
562                message: "boom".to_string(),
563            },
564            LlmResponseChunk::DecodeError {
565                message: "boom".to_string(),
566            },
567        ] {
568            for target in [
569                WireFormat::OpenAiChat,
570                WireFormat::OpenAiResponses,
571                WireFormat::AnthropicMessages,
572            ] {
573                let chunks: LlmResponseStream = stream::iter(vec![
574                    Ok(LlmResponseChunk::TextDelta {
575                        index: 0,
576                        text: "before".to_string(),
577                    }
578                    .into()),
579                    Ok(message.clone().into()),
580                    Ok(LlmResponseChunk::TextDelta {
581                        index: 0,
582                        text: "after".to_string(),
583                    }
584                    .into()),
585                ])
586                .boxed();
587                let results = block_on(encode_stream(chunks, target, None)?.collect::<Vec<_>>());
588                let (data, errors): (Vec<_>, Vec<_>) = results.into_iter().partition(Result::is_ok);
589                let data = data.into_iter().flatten().collect::<Vec<Value>>();
590                let body = serde_json::to_string(&data)?;
591
592                assert!(
593                    body.contains("before"),
594                    "{target:?}: pre-error content missing:\n{body}"
595                );
596                assert!(
597                    !body.contains("after"),
598                    "{target:?}/{message:?}: content leaked after the error:\n{body}"
599                );
600                // The error rides the error channel, carrying the codec's own encoded event
601                // so the consumer can forward the upstream payload rather than rebuild it.
602                let [Err(LlmStreamError::Upstream(event))] = errors.as_slice() else {
603                    panic!("{target:?}/{message:?}: expected one in-band error, got {errors:?}");
604                };
605                assert!(
606                    serde_json::to_string(event)?.contains("boom"),
607                    "{target:?}: error event lost its message: {event}"
608                );
609            }
610        }
611        Ok(())
612    }
613
614    // A clean stop stays entirely on the data channel, so a consumer that appends a
615    // format-level success sentinel only does so for a turn that actually completed.
616    #[test]
617    fn encode_stream_keeps_a_normal_stop_off_the_error_channel() -> Result<(), BoxError> {
618        let chunks: LlmResponseStream = stream::iter(vec![
619            Ok(LlmResponseChunk::TextDelta {
620                index: 0,
621                text: "hi".to_string(),
622            }
623            .into()),
624            Ok(LlmResponseChunk::MessageStop {
625                reason: Some("stop".to_string()),
626            }
627            .into()),
628        ])
629        .boxed();
630
631        let results =
632            block_on(encode_stream(chunks, WireFormat::OpenAiChat, None)?.collect::<Vec<_>>());
633
634        assert!(
635            results.iter().all(Result::is_ok),
636            "a completed turn reported an error: {results:?}"
637        );
638        Ok(())
639    }
640
641    // A replayed provider error ends the stream before the encoder polls the source again.
642    #[test]
643    fn encode_stream_stops_polling_after_a_replayed_error() -> Result<(), BoxError> {
644        let error = LlmResponseStreamEvent::preserved(
645            WireFormat::OpenAiResponses,
646            json!({"type": "error", "message": "boom"}),
647            vec![LlmResponseChunk::StreamError {
648                message: "boom".to_string(),
649            }],
650        );
651        let chunks: LlmResponseStream = stream::iter([Ok(error)])
652            .chain(stream::poll_fn(|_| {
653                panic!("encode_stream polled the source after an in-band error")
654            }))
655            .boxed();
656
657        let events =
658            block_on(encode_stream(chunks, WireFormat::OpenAiResponses, None)?.collect::<Vec<_>>());
659
660        // The replayed provider event is the terminal frame, so it arrives verbatim on the
661        // error channel and nothing follows it.
662        let [Err(LlmStreamError::Upstream(event))] = events.as_slice() else {
663            return Err(format!("expected one in-band error, got {events:?}").into());
664        };
665        assert_eq!(*event, json!({"type": "error", "message": "boom"}));
666        Ok(())
667    }
668
669    // The guard keys on `errored`, not `finished`, so a normal completion still emits the
670    // trailing usage chunk the OpenAI chat codec reports only after `finished` is set.
671    #[test]
672    fn encode_stream_keeps_trailing_usage_after_a_normal_stop() -> Result<(), BoxError> {
673        let chunks: LlmResponseStream = stream::iter(vec![
674            Ok(LlmResponseChunk::TextDelta {
675                index: 0,
676                text: "hi".to_string(),
677            }
678            .into()),
679            Ok(LlmResponseChunk::MessageStop {
680                reason: Some("stop".to_string()),
681            }
682            .into()),
683            Ok(LlmResponseChunk::Usage(switchyard_protocol::llm::Usage {
684                output_tokens: Some(7),
685                ..Default::default()
686            })
687            .into()),
688        ])
689        .boxed();
690        let events =
691            block_on(encode_stream(chunks, WireFormat::OpenAiChat, None)?.collect::<Vec<_>>())
692                .into_iter()
693                .collect::<Result<Vec<Value>, LlmStreamError>>()?;
694        let body = serde_json::to_string(&events)?;
695        assert!(
696            events
697                .iter()
698                .any(|event| event["choices"][0]["finish_reason"] == "stop"),
699            "missing stop terminal:\n{body}"
700        );
701        assert!(
702            body.contains("\"usage\""),
703            "trailing usage dropped after a normal stop:\n{body}"
704        );
705        Ok(())
706    }
707
708    #[test]
709    fn decode_stream_parses_sse_bytes_into_ir_chunks() -> Result<(), LlmClientError> {
710        let sse = b"data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\
711             data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n\
712             data: [DONE]\n\n\
713             data: {\"choices\":[{\"delta\":{\"content\":\" ignored\"}}]}\n\n"
714            .to_vec();
715        let bytes = stream::once(async move { Ok::<Vec<u8>, LlmClientError>(sse) });
716        let chunks = decode_all(bytes, WireFormat::OpenAiChat)?;
717        assert_eq!(text_of(&chunks), "Hello world");
718        Ok(())
719    }
720
721    #[test]
722    fn stream_helpers_replay_same_format_provider_fields() -> Result<(), BoxError> {
723        let provider_event = json!({
724            "id": "chatcmpl-test",
725            "object": "chat.completion.chunk",
726            "system_fingerprint": "fp_provider_specific",
727            "choices": [{
728                "index": 0,
729                "delta": {"content": "Hello"},
730                "finish_reason": "stop"
731            }]
732        });
733        let bytes = stream::once({
734            let frame = format!("data: {provider_event}\n\n").into_bytes();
735            async move { Ok::<Vec<u8>, LlmClientError>(frame) }
736        });
737        let decoded = decode_stream(bytes, WireFormat::OpenAiChat)?;
738        let replayed =
739            block_on(encode_stream(decoded, WireFormat::OpenAiChat, None)?.collect::<Vec<_>>())
740                .into_iter()
741                .collect::<Result<Vec<Value>, LlmStreamError>>()?;
742
743        assert_eq!(replayed, vec![provider_event]);
744        Ok(())
745    }
746
747    #[test]
748    fn openai_chat_replay_stamps_the_served_model_without_losing_extensions() {
749        let mut event = json!({
750            "choices": [{"delta": {"content": "Hello"}}],
751            "system_fingerprint": "fp_provider_specific",
752        });
753
754        stamp_streamed_response_model(&mut event, WireFormat::OpenAiChat, Some("served/model"));
755
756        assert_eq!(event["model"], "served/model");
757        assert_eq!(event["system_fingerprint"], "fp_provider_specific");
758    }
759
760    #[test]
761    fn responses_replay_stamps_the_served_model_inside_response() {
762        let mut event = json!({
763            "type": "response.created",
764            "response": {
765                "id": "resp_1",
766                "model": "provider/model",
767                "provider_extension": true,
768            },
769        });
770
771        stamp_streamed_response_model(
772            &mut event,
773            WireFormat::OpenAiResponses,
774            Some("served/model"),
775        );
776
777        assert_eq!(event["response"]["model"], "served/model");
778        assert_eq!(event["response"]["provider_extension"], true);
779    }
780
781    #[test]
782    fn anthropic_replay_stamps_the_served_model_inside_message() {
783        let mut event = json!({
784            "type": "message_start",
785            "message": {
786                "id": "msg_1",
787                "model": "provider/model",
788                "provider_extension": true,
789            },
790        });
791
792        stamp_streamed_response_model(
793            &mut event,
794            WireFormat::AnthropicMessages,
795            Some("served/model"),
796        );
797
798        assert_eq!(event["message"]["model"], "served/model");
799        assert_eq!(event["message"]["provider_extension"], true);
800    }
801
802    #[test]
803    fn decode_stream_reassembles_frames_split_across_chunks() -> Result<(), BoxError> {
804        // A multi-byte codepoint and the frame boundaries are split across
805        // one-byte chunks; the BufReader must reassemble them losslessly.
806        let payload = json!({"choices": [{"delta": {"content": "café"}}]});
807        let sse = format!("data: {payload}\n\ndata: [DONE]\n\n");
808        let bytes = stream::iter(
809            sse.into_bytes()
810                .into_iter()
811                .map(|byte| Ok::<Vec<u8>, LlmClientError>(vec![byte])),
812        );
813        let chunks = decode_all(bytes, WireFormat::OpenAiChat)?;
814        assert_eq!(text_of(&chunks), "café");
815        Ok(())
816    }
817
818    #[test]
819    fn decode_stream_decodes_trailing_frame_without_blank_line() -> Result<(), BoxError> {
820        // A non-standard upstream omits the final blank line; the last frame
821        // must still be decoded rather than dropped.
822        let sse =
823            b"data: {\"choices\":[{\"delta\":{\"content\":\"tail\"}}]}\n\ndata: [DONE]".to_vec();
824        let bytes = stream::once(async move { Ok::<Vec<u8>, LlmClientError>(sse) });
825        let chunks = decode_all(bytes, WireFormat::OpenAiChat)?;
826        assert_eq!(text_of(&chunks), "tail");
827        Ok(())
828    }
829
830    #[test]
831    fn decode_stream_rejects_eof_before_a_terminal_event() -> Result<(), BoxError> {
832        // Preserve valid content before surfacing premature EOF as the terminal stream error.
833        let sse = b"data: {\"choices\":[{\"delta\":{\"content\":\"partial\"},\"finish_reason\":null}]}\n\n".to_vec();
834        let bytes = stream::once(async move { Ok::<Vec<u8>, LlmClientError>(sse) });
835        let results = block_on(decode_stream(bytes, WireFormat::OpenAiChat)?.collect::<Vec<_>>());
836
837        let Some(Ok(first)) = results.first() else {
838            return Err("expected the partial event".into());
839        };
840        assert_eq!(text_of(std::slice::from_ref(first)), "partial");
841        let Some(Err(LlmClientError::ResponseTranslation(message))) = results.last() else {
842            panic!("expected incomplete OpenAI stream to fail");
843        };
844        assert_eq!(message, "openai_chat stream ended before a terminal event");
845        Ok(())
846    }
847
848    #[test]
849    fn responses_failed_before_done_remains_a_stream_error() -> Result<(), BoxError> {
850        let sse = b"data: {\"type\":\"response.failed\",\"response\":{\"error\":{\"message\":\"upstream failed\"}}}\n\ndata: [DONE]\n\n".to_vec();
851        let bytes = stream::once(async move { Ok::<Vec<u8>, LlmClientError>(sse) });
852        let events = decode_all(bytes, WireFormat::OpenAiResponses)?;
853
854        let Some(LlmResponseChunk::StreamError { message }) =
855            events.first().and_then(|event| event.normalized().first())
856        else {
857            return Err("expected response.failed to decode as StreamError".into());
858        };
859        assert_eq!(message, "upstream failed");
860
861        let chunks: LlmResponseStream = stream::iter(
862            events
863                .into_iter()
864                .map(Ok::<LlmResponseStreamEvent, LlmClientError>),
865        )
866        .boxed();
867        let replayed =
868            block_on(encode_stream(chunks, WireFormat::OpenAiResponses, None)?.collect::<Vec<_>>());
869
870        // `response.failed` stays a terminal error, replayed with its payload intact.
871        let [Err(LlmStreamError::Upstream(event))] = replayed.as_slice() else {
872            return Err(format!("expected one in-band error, got {replayed:?}").into());
873        };
874        assert_eq!(
875            *event,
876            json!({
877                "type": "response.failed",
878                "response": {"error": {"message": "upstream failed"}}
879            })
880        );
881        Ok(())
882    }
883
884    #[test]
885    fn decode_stream_decodes_crlf_delimited_frames() -> Result<(), BoxError> {
886        // CRLF framing: blank lines are `\r\n\r\n` and the bare `\r` must not
887        // block the frame boundary.
888        let sse =
889            b"data: {\"choices\":[{\"delta\":{\"content\":\"crlf\"}}]}\r\n\r\ndata: [DONE]\r\n\r\n"
890                .to_vec();
891        let bytes = stream::once(async move { Ok::<Vec<u8>, LlmClientError>(sse) });
892        let chunks = decode_all(bytes, WireFormat::OpenAiChat)?;
893        assert_eq!(text_of(&chunks), "crlf");
894        Ok(())
895    }
896
897    #[test]
898    fn decode_stream_propagates_source_errors() -> Result<(), BoxError> {
899        // A transport error mid-stream surfaces as an error item, not a panic.
900        let bytes = stream::iter(vec![
901            Ok::<Vec<u8>, LlmClientError>(
902                b"data: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n".to_vec(),
903            ),
904            Err::<Vec<u8>, LlmClientError>(LlmClientError::Transport {
905                source: Box::new(std::io::Error::other("upstream exploded")),
906            }),
907        ]);
908        let results = block_on(decode_stream(bytes, WireFormat::OpenAiChat)?.collect::<Vec<_>>());
909        let Some(Err(error)) = results.last() else {
910            panic!("expected the source error");
911        };
912        assert!(matches!(error, LlmClientError::Transport { .. }));
913        Ok(())
914    }
915
916    #[test]
917    fn decode_stream_classifies_invalid_sse_json() -> Result<(), BoxError> {
918        let bytes =
919            stream::once(async { Ok::<Vec<u8>, LlmClientError>(b"data: {invalid}\n\n".to_vec()) });
920        let results = block_on(decode_stream(bytes, WireFormat::OpenAiChat)?.collect::<Vec<_>>());
921        let Some(Err(error)) = results.last() else {
922            panic!("expected invalid SSE JSON to fail");
923        };
924        assert!(matches!(error, LlmClientError::ResponseTranslation(_)));
925        Ok(())
926    }
927}
switchyard_translation/helpers.rs

switchyard_translation/
helpers.rs
