switchyard_translation/codecs/
stream.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Shared streaming codec contracts, registry, and stream state.
5
6use std::collections::BTreeMap;
7use std::sync::Arc;
8
9use serde::{Deserialize, Serialize};
10use serde_json::{Map, Value, json};
11
12use crate::LlmResponseChunk;
13use crate::codecs::anthropic::AnthropicMessagesStreamCodec;
14use crate::codecs::openai_chat::OpenAiChatStreamCodec;
15use crate::codecs::responses::OpenAiResponsesStreamCodec;
16use crate::engine::{FormatRegistry, TranslationEngine};
17use crate::error::{Result, TranslationError};
18use crate::format::{FormatId, WireFormat};
19use crate::llm::Usage;
20
21/// Mutable state accumulated while translating one streaming response.
22#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
23pub struct StreamTranslationState {
24    pub source: Option<FormatId>,
25    pub target: Option<FormatId>,
26    /// Model name observed on the source provider stream.
27    pub model: Option<String>,
28    /// Message/response ID observed on the source provider stream.
29    pub message_id: Option<String>,
30    /// Model that served the call, exposed to the client in place of the id the
31    /// source stream reported. `None` falls back to [`Self::model`].
32    pub target_model: Option<String>,
33    /// Optional message/response ID the target stream should expose to the client.
34    pub target_message_id: Option<String>,
35    pub saw_message_start: bool,
36    pub emitted_message_start: bool,
37    pub finished: bool,
38    /// Set once an in-band error event was emitted; the encoder then emits nothing further.
39    pub errored: bool,
40    pub usage: Usage,
41
42    pub(crate) output_tokens_seen: u64,
43    pub(crate) saw_backend_usage: bool,
44    pub(crate) stop_reason: Option<String>,
45    /// Refusal metadata carried by the source, replayed instead of a synthesized
46    /// object so a named policy category and its explanation are not flattened away.
47    pub(crate) stop_details: Option<Value>,
48    pub(crate) emitted_message_delta: bool,
49
50    pub(crate) next_content_index: usize,
51    pub(crate) text_block_index: Option<usize>,
52    pub(crate) text_block_started: bool,
53    pub(crate) emitted_content_block: bool,
54    pub(crate) tool_states: BTreeMap<usize, StreamToolState>,
55    #[serde(default)]
56    pub(crate) pending_chat_tool_names: BTreeMap<usize, String>,
57    #[serde(default)]
58    pub(crate) active_anthropic_tool: Option<usize>,
59    #[serde(default)]
60    pub(crate) deferred_anthropic_tools: Vec<usize>,
61    /// Empty initial inputs remain placeholders until their blocks close without arguments.
62    #[serde(default)]
63    pub(crate) empty_anthropic_tool_inputs: std::collections::BTreeSet<usize>,
64    /// Reasoning text observed while DECODING, per output index, so a completed item
65    /// that repeats already-streamed text is not decoded twice.
66    pub(crate) decoded_reasoning: BTreeMap<usize, String>,
67    /// Output indexes whose encrypted reasoning payload was already decoded.
68    pub(crate) decoded_reasoning_encrypted: std::collections::BTreeSet<usize>,
69    /// Set once a tool call was observed while DECODING, so a terminal event that names no
70    /// stop reason can still report tool use.
71    pub(crate) decoded_tool_call: bool,
72    /// Text decoded from Responses, keyed by output index and then content index.
73    #[serde(default)]
74    pub(crate) decoded_response_text: BTreeMap<usize, BTreeMap<usize, String>>,
75
76    pub(crate) response_created: bool,
77    pub(crate) response_text_started: bool,
78    pub(crate) response_text_output_index: Option<usize>,
79    pub(crate) response_text: String,
80    /// Reasoning items being ENCODED, keyed by source content index. A response can carry
81    /// several reasoning items (GPT-5 emits one ahead of each tool call); each becomes its
82    /// own output item so none of their encrypted payloads is lost.
83    pub(crate) response_reasoning: BTreeMap<usize, ResponseReasoningState>,
84    /// Source indexes whose provider reasoning item id was already announced while DECODING.
85    pub(crate) decoded_reasoning_ids: std::collections::BTreeSet<usize>,
86    pub(crate) next_response_output_index: usize,
87    pub(crate) response_sequence_number: u64,
88    pub(crate) next_chat_tool_index: usize,
89
90    pub(crate) reasoning_block_index: Option<usize>,
91    pub(crate) reasoning_block_started: bool,
92}
93
94// One Responses reasoning output item under construction by the encoder.
95#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
96pub(crate) struct ResponseReasoningState {
97    /// Set once the item's `added` event was emitted.
98    pub(crate) started: bool,
99    /// Set once the item's summary part opened, which happens on the first text delta. An
100    /// encrypted-only item never opens one, so it never has to close one either.
101    pub(crate) summary_started: bool,
102    pub(crate) output_index: Option<usize>,
103    /// Provider item id the encrypted reasoning was issued under. Used as the emitted item id
104    /// so the client's replay verifies upstream; `None` falls back to a synthesized id.
105    pub(crate) item_id: Option<String>,
106    pub(crate) text: String,
107    /// Opaque `encrypted_content` carried by a Responses reasoning item. Kept verbatim so
108    /// the emitted item stays replayable by the client even when no plaintext streamed.
109    pub(crate) encrypted: Option<String>,
110    #[serde(default)]
111    pub(crate) anthropic_signature: Option<String>,
112    /// Summary text held back while the item's provider id is still unknown. The item opens
113    /// once that id arrives, or once later output shows none is coming.
114    #[serde(default)]
115    pub(crate) pending_text: String,
116}
117
118// Tracks an in-progress streamed tool call across provider-specific deltas.
119#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
120pub(crate) struct StreamToolState {
121    pub(crate) id: Option<String>,
122    pub(crate) name: Option<String>,
123    pub(crate) arguments: String,
124    /// Arguments observed while DECODING the source stream.
125    ///
126    /// Separate from `arguments`, which encoders accumulate. A decoder that
127    /// deduplicates against `arguments` only works when one state performs
128    /// both halves of the translation; when a caller buffers a stream with its
129    /// own state and encodes later, the field is empty and the duplicate is
130    /// emitted.
131    pub(crate) decoded_arguments: String,
132    #[serde(default)]
133    pub(crate) has_decoded_identity: bool,
134    /// True when the source announced a `custom_tool_call`. The decoder then reads the call's
135    /// input only from the completed item, because some providers also stream argument deltas
136    /// whose JSON does not match that input.
137    #[serde(default)]
138    pub(crate) decoded_custom_call: bool,
139    pub(crate) pending_arguments: String,
140    pub(crate) started: bool,
141    pub(crate) content_index: Option<usize>,
142    /// Position in the OpenAI Chat `tool_calls` array, assigned when ENCODING. Chat numbers
143    /// tool calls on their own, while Anthropic and Responses index the whole content array,
144    /// so the source index cannot be reused.
145    pub(crate) chat_tool_index: Option<usize>,
146    pub(crate) response_output_index: Option<usize>,
147    pub(crate) response_item_id: Option<String>,
148}
149
150impl StreamTranslationState {
151    /// Creates stream state with source and target formats already attached.
152    pub fn new(source: impl Into<FormatId>, target: impl Into<FormatId>) -> Self {
153        Self {
154            source: Some(source.into()),
155            target: Some(target.into()),
156            ..Self::default()
157        }
158    }
159}
160
161/// Registry-backed streaming translator.
162#[derive(Default)]
163pub struct StreamTranslationEngine {
164    engine: TranslationEngine,
165}
166
167/// Codec contract for one provider streaming event format.
168pub trait StreamCodec: Send + Sync {
169    /// Returns the stream format handled by this codec.
170    fn format(&self) -> FormatId;
171
172    /// Returns whether a JSON event establishes stream completion.
173    fn is_terminal_event(&self, event: &Value) -> bool;
174
175    /// Returns whether the SSE `[DONE]` marker establishes stream completion.
176    ///
177    /// The marker ends consumption even when this returns `false`; a terminal JSON event
178    /// must then have been observed to avoid a premature-EOF error.
179    fn sse_done_is_terminal(&self) -> bool;
180
181    /// Overrides the model reported by an encoded or preserved event.
182    ///
183    /// Called only when a served model was supplied, including for events emitted by `finish`.
184    fn stamp_streamed_response_model(&self, event: &mut Value, served_model: &str);
185
186    /// Decodes one provider event into zero or more neutral events.
187    fn decode_event(
188        &self,
189        state: &mut StreamTranslationState,
190        event: &Value,
191    ) -> Vec<LlmResponseChunk>;
192
193    /// Encodes one neutral event into zero or more provider events.
194    fn encode_event(
195        &self,
196        state: &mut StreamTranslationState,
197        event: LlmResponseChunk,
198    ) -> Vec<Value>;
199
200    /// Advances encoder state after an exact same-format event replay.
201    ///
202    /// Exact replay returns the preserved provider JSON instead of the JSON emitted by
203    /// [`Self::encode_event`]. The encoder must nevertheless observe the normalized chunks so
204    /// [`Self::finish`] can close an incomplete stream without duplicating an already replayed
205    /// terminal event. Codecs whose terminal state cannot be inferred from `MessageStop` alone
206    /// may override this hook and inspect `raw`.
207    fn observe_replayed_event(
208        &self,
209        state: &mut StreamTranslationState,
210        _raw: &Value,
211        normalized: Vec<LlmResponseChunk>,
212    ) {
213        let replayed_terminal = normalized
214            .iter()
215            .any(|chunk| matches!(chunk, LlmResponseChunk::MessageStop { .. }));
216        for chunk in normalized {
217            drop(self.encode_event(state, chunk));
218        }
219        if replayed_terminal {
220            state.finished = true;
221        }
222    }
223
224    /// Emits any terminal provider events needed after the source stream ends.
225    ///
226    /// This is intentionally required on every codec. Some target formats
227    /// need explicit terminal events after the source closes (for example,
228    /// Anthropic ``message_delta``/``message_stop`` or Responses
229    /// ``response.completed``). Formats that have no source-close work should
230    /// return an empty vector explicitly so the no-op behavior is a conscious
231    /// codec-level choice.
232    fn finish(&self, state: &mut StreamTranslationState) -> Vec<Value>;
233}
234
235/// Registry mapping stream wire formats to stream codecs.
236#[derive(Default)]
237pub struct StreamCodecRegistry {
238    codecs: BTreeMap<FormatId, Arc<dyn StreamCodec>>,
239}
240
241impl StreamCodecRegistry {
242    /// Creates an empty stream codec registry.
243    pub fn new() -> Self {
244        Self::default()
245    }
246
247    /// Creates a registry populated with built-in stream codecs.
248    pub fn with_builtins() -> Self {
249        let mut registry = Self::new();
250        registry.register(OpenAiChatStreamCodec);
251        registry.register(AnthropicMessagesStreamCodec);
252        registry.register(OpenAiResponsesStreamCodec);
253        registry
254    }
255
256    /// Registers or replaces a stream codec for its declared format.
257    pub fn register(&mut self, codec: impl StreamCodec + 'static) {
258        self.codecs.insert(codec.format(), Arc::new(codec));
259    }
260
261    /// Looks up a stream codec by format identifier.
262    pub fn codec(&self, format: impl Into<FormatId>) -> Result<Arc<dyn StreamCodec>> {
263        let format = format.into();
264        self.codecs.get(&format).cloned().ok_or_else(|| {
265            TranslationError::Other(format!("no stream codec registered for {format}"))
266        })
267    }
268}
269
270impl StreamTranslationEngine {
271    /// Creates a streaming engine from an explicit codec registry.
272    pub fn new(registry: StreamCodecRegistry) -> Self {
273        Self {
274            engine: TranslationEngine::with_registries(FormatRegistry::with_builtins(), registry),
275        }
276    }
277
278    /// Translates one source provider event into target provider events.
279    pub fn translate_event(
280        &self,
281        state: &mut StreamTranslationState,
282        source: impl Into<FormatId>,
283        target: impl Into<FormatId>,
284        event: &Value,
285    ) -> Result<Vec<Value>> {
286        self.engine.translate_event(state, source, target, event)
287    }
288
289    /// Finishes target-provider stream emission after the source stream closes.
290    pub fn finish(
291        &self,
292        state: &mut StreamTranslationState,
293        target: impl Into<FormatId>,
294    ) -> Result<Vec<Value>> {
295        self.engine.finish_stream(state, target)
296    }
297
298    /// Convenience helper using built-in codecs and error events instead of `Result`.
299    pub fn translate_event_with_builtins(
300        state: &mut StreamTranslationState,
301        source: impl Into<FormatId>,
302        target: impl Into<FormatId>,
303        event: &Value,
304    ) -> Vec<Value> {
305        Self::default()
306            .translate_event(state, source, target, event)
307            .unwrap_or_else(|error| vec![json!({"error": {"message": error.to_string()}})])
308    }
309}
310
311/// Decodes one provider stream event with the built-in codec registry.
312pub fn decode_stream_event(
313    state: &mut StreamTranslationState,
314    source: impl Into<FormatId>,
315    event: &Value,
316) -> Vec<LlmResponseChunk> {
317    let source = source.into();
318    StreamCodecRegistry::with_builtins()
319        .codec(source)
320        .map(|codec| codec.decode_event(state, event))
321        .unwrap_or_else(|error| {
322            // A missing codec is a translation-side failure, not something the
323            // upstream sent, so it decodes to `DecodeError` rather than `StreamError`.
324            vec![LlmResponseChunk::DecodeError {
325                message: error.to_string(),
326            }]
327        })
328}
329
330pub(crate) fn encode_response_stream_event(
331    state: &mut StreamTranslationState,
332    target_codec: &dyn StreamCodec,
333    target: &FormatId,
334    event: crate::LlmResponseStreamEvent,
335) -> Vec<Value> {
336    if state.errored {
337        return Vec::new();
338    }
339    let (preservation, normalized) = event.into_parts();
340    if let Some(preservation) = preservation {
341        let (source, raw) = preservation.into_parts();
342        if let Err(error) = super::responses::validate_stream_output(&source, target, &raw) {
343            return target_codec.encode_event(
344                state,
345                LlmResponseChunk::DecodeError {
346                    message: error.to_string(),
347                },
348            );
349        }
350        // Invalid protocol data must become an error frame, not be replayed as ordinary data.
351        let has_decode_error = normalized
352            .iter()
353            .any(|chunk| matches!(chunk, LlmResponseChunk::DecodeError { .. }));
354        if &source == target && !has_decode_error {
355            // Exact replay bypasses the target encoder's emitted JSON, but the encoder must
356            // still observe every normalized chunk. Otherwise `finish` starts from empty state:
357            // a clean EOF after a nonterminal provider event can omit or synthesize malformed
358            // terminal events, while a replayed terminal can be emitted twice.
359            target_codec.observe_replayed_event(state, &raw, normalized);
360            return vec![raw];
361        }
362    }
363
364    normalized
365        .into_iter()
366        .flat_map(|chunk| target_codec.encode_event(state, chunk))
367        .collect()
368}
369
370/// Encodes one neutral stream event with the built-in codec registry.
371pub fn encode_stream_event(
372    state: &mut StreamTranslationState,
373    target: impl Into<FormatId>,
374    event: LlmResponseChunk,
375) -> Vec<Value> {
376    StreamCodecRegistry::with_builtins()
377        .codec(target)
378        .map(|codec| codec.encode_event(state, event))
379        .unwrap_or_else(|error| vec![json!({"error": {"message": error.to_string()}})])
380}
381
382// Records source-provider identity carried by decoded stream events.
383pub(crate) fn record_source_identity(
384    state: &mut StreamTranslationState,
385    id: Option<String>,
386    model: Option<String>,
387) {
388    if id.is_some() {
389        state.message_id = id;
390    }
391    if model.is_some() {
392        state.model = model;
393    }
394}
395
396// Returns the served model when supplied, otherwise the upstream model.
397pub(crate) fn target_model_or_source_model(state: &StreamTranslationState) -> String {
398    state
399        .target_model
400        .clone()
401        .or_else(|| state.model.clone())
402        .unwrap_or_else(|| "unknown".to_string())
403}
404
405// Returns the target/client ID when supplied, otherwise the upstream ID.
406pub(crate) fn target_message_id_or_source_message_id(
407    state: &StreamTranslationState,
408) -> Option<&str> {
409    state
410        .target_message_id
411        .as_deref()
412        .or(state.message_id.as_deref())
413}
414
415// Checks whether the current source format matches a built-in format.
416pub(crate) fn state_source_is(state: &StreamTranslationState, format: WireFormat) -> bool {
417    let format_id: FormatId = format.into();
418    match &state.source {
419        Some(source) => source == &format_id,
420        None => false,
421    }
422}
423
424// Reads a non-empty string field from an event object.
425pub(crate) fn string_field(object: &Map<String, Value>, key: &str) -> Option<String> {
426    object
427        .get(key)
428        .and_then(Value::as_str)
429        .filter(|value| !value.is_empty())
430        .map(ToOwned::to_owned)
431}
switchyard_translation/codecs/stream.rs

switchyard_translation/codecs/
stream.rs
