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 decoding stops after yielding a terminal JSON event.
176    ///
177    /// The default keeps reading so codecs can emit trailing events, such as usage.
178    fn stops_at_terminal_event(&self) -> bool {
179        false
180    }
181
182    /// Returns whether the SSE `[DONE]` marker establishes stream completion.
183    ///
184    /// The marker ends consumption even when this returns `false`; a terminal JSON event
185    /// must then have been observed to avoid a premature-EOF error.
186    fn sse_done_is_terminal(&self) -> bool;
187
188    /// Overrides the model reported by an encoded or preserved event.
189    ///
190    /// Called only when a served model was supplied, including for events emitted by `finish`.
191    fn stamp_streamed_response_model(&self, event: &mut Value, served_model: &str);
192
193    /// Decodes one provider event into zero or more neutral events.
194    fn decode_event(
195        &self,
196        state: &mut StreamTranslationState,
197        event: &Value,
198    ) -> Vec<LlmResponseChunk>;
199
200    /// Encodes one neutral event into zero or more provider events.
201    fn encode_event(
202        &self,
203        state: &mut StreamTranslationState,
204        event: LlmResponseChunk,
205    ) -> Vec<Value>;
206
207    /// Advances encoder state after an exact same-format event replay.
208    ///
209    /// Exact replay returns the preserved provider JSON instead of the JSON emitted by
210    /// [`Self::encode_event`]. The encoder must nevertheless observe the normalized chunks so
211    /// [`Self::finish`] can close an incomplete stream without duplicating an already replayed
212    /// terminal event. Codecs whose terminal state cannot be inferred from `MessageStop` alone
213    /// may override this hook and inspect `raw`.
214    fn observe_replayed_event(
215        &self,
216        state: &mut StreamTranslationState,
217        _raw: &Value,
218        normalized: Vec<LlmResponseChunk>,
219    ) {
220        let replayed_terminal = normalized
221            .iter()
222            .any(|chunk| matches!(chunk, LlmResponseChunk::MessageStop { .. }));
223        for chunk in normalized {
224            drop(self.encode_event(state, chunk));
225        }
226        if replayed_terminal {
227            state.finished = true;
228        }
229    }
230
231    /// Emits any terminal provider events needed after the source stream ends.
232    ///
233    /// This is intentionally required on every codec. Some target formats
234    /// need explicit terminal events after the source closes (for example,
235    /// Anthropic ``message_delta``/``message_stop`` or Responses
236    /// ``response.completed``). Formats that have no source-close work should
237    /// return an empty vector explicitly so the no-op behavior is a conscious
238    /// codec-level choice.
239    fn finish(&self, state: &mut StreamTranslationState) -> Vec<Value>;
240}
241
242/// Registry mapping stream wire formats to stream codecs.
243#[derive(Default)]
244pub struct StreamCodecRegistry {
245    codecs: BTreeMap<FormatId, Arc<dyn StreamCodec>>,
246}
247
248impl StreamCodecRegistry {
249    /// Creates an empty stream codec registry.
250    pub fn new() -> Self {
251        Self::default()
252    }
253
254    /// Creates a registry populated with built-in stream codecs.
255    pub fn with_builtins() -> Self {
256        let mut registry = Self::new();
257        registry.register(OpenAiChatStreamCodec);
258        registry.register(AnthropicMessagesStreamCodec);
259        registry.register(OpenAiResponsesStreamCodec);
260        registry
261    }
262
263    /// Registers or replaces a stream codec for its declared format.
264    pub fn register(&mut self, codec: impl StreamCodec + 'static) {
265        self.codecs.insert(codec.format(), Arc::new(codec));
266    }
267
268    /// Looks up a stream codec by format identifier.
269    pub fn codec(&self, format: impl Into<FormatId>) -> Result<Arc<dyn StreamCodec>> {
270        let format = format.into();
271        self.codecs.get(&format).cloned().ok_or_else(|| {
272            TranslationError::Other(format!("no stream codec registered for {format}"))
273        })
274    }
275}
276
277impl StreamTranslationEngine {
278    /// Creates a streaming engine from an explicit codec registry.
279    pub fn new(registry: StreamCodecRegistry) -> Self {
280        Self {
281            engine: TranslationEngine::with_registries(FormatRegistry::with_builtins(), registry),
282        }
283    }
284
285    /// Translates one source provider event into target provider events.
286    pub fn translate_event(
287        &self,
288        state: &mut StreamTranslationState,
289        source: impl Into<FormatId>,
290        target: impl Into<FormatId>,
291        event: &Value,
292    ) -> Result<Vec<Value>> {
293        self.engine.translate_event(state, source, target, event)
294    }
295
296    /// Finishes target-provider stream emission after the source stream closes.
297    pub fn finish(
298        &self,
299        state: &mut StreamTranslationState,
300        target: impl Into<FormatId>,
301    ) -> Result<Vec<Value>> {
302        self.engine.finish_stream(state, target)
303    }
304
305    /// Convenience helper using built-in codecs and error events instead of `Result`.
306    pub fn translate_event_with_builtins(
307        state: &mut StreamTranslationState,
308        source: impl Into<FormatId>,
309        target: impl Into<FormatId>,
310        event: &Value,
311    ) -> Vec<Value> {
312        Self::default()
313            .translate_event(state, source, target, event)
314            .unwrap_or_else(|error| vec![json!({"error": {"message": error.to_string()}})])
315    }
316}
317
318/// Decodes one provider stream event with the built-in codec registry.
319pub fn decode_stream_event(
320    state: &mut StreamTranslationState,
321    source: impl Into<FormatId>,
322    event: &Value,
323) -> Vec<LlmResponseChunk> {
324    let source = source.into();
325    StreamCodecRegistry::with_builtins()
326        .codec(source)
327        .map(|codec| codec.decode_event(state, event))
328        .unwrap_or_else(|error| {
329            // A missing codec is a translation-side failure, not something the
330            // upstream sent, so it decodes to `DecodeError` rather than `StreamError`.
331            vec![LlmResponseChunk::DecodeError {
332                message: error.to_string(),
333            }]
334        })
335}
336
337pub(crate) fn encode_response_stream_event(
338    state: &mut StreamTranslationState,
339    target_codec: &dyn StreamCodec,
340    target: &FormatId,
341    event: crate::LlmResponseStreamEvent,
342) -> Vec<Value> {
343    if state.errored {
344        return Vec::new();
345    }
346    let (preservation, normalized) = event.into_parts();
347    if let Some(preservation) = preservation {
348        let (source, raw) = preservation.into_parts();
349        if let Err(error) = super::responses::validate_stream_output(&source, target, &raw) {
350            return target_codec.encode_event(
351                state,
352                LlmResponseChunk::DecodeError {
353                    message: error.to_string(),
354                },
355            );
356        }
357        // Invalid protocol data must become an error frame, not be replayed as ordinary data.
358        let has_decode_error = normalized
359            .iter()
360            .any(|chunk| matches!(chunk, LlmResponseChunk::DecodeError { .. }));
361        if &source == target && !has_decode_error {
362            // Exact replay bypasses the target encoder's emitted JSON, but the encoder must
363            // still observe every normalized chunk. Otherwise `finish` starts from empty state:
364            // a clean EOF after a nonterminal provider event can omit or synthesize malformed
365            // terminal events, while a replayed terminal can be emitted twice.
366            target_codec.observe_replayed_event(state, &raw, normalized);
367            return vec![raw];
368        }
369    }
370
371    normalized
372        .into_iter()
373        .flat_map(|chunk| target_codec.encode_event(state, chunk))
374        .collect()
375}
376
377/// Encodes one neutral stream event with the built-in codec registry.
378pub fn encode_stream_event(
379    state: &mut StreamTranslationState,
380    target: impl Into<FormatId>,
381    event: LlmResponseChunk,
382) -> Vec<Value> {
383    StreamCodecRegistry::with_builtins()
384        .codec(target)
385        .map(|codec| codec.encode_event(state, event))
386        .unwrap_or_else(|error| vec![json!({"error": {"message": error.to_string()}})])
387}
388
389// Records source-provider identity carried by decoded stream events.
390pub(crate) fn record_source_identity(
391    state: &mut StreamTranslationState,
392    id: Option<String>,
393    model: Option<String>,
394) {
395    if id.is_some() {
396        state.message_id = id;
397    }
398    if model.is_some() {
399        state.model = model;
400    }
401}
402
403// Returns the served model when supplied, otherwise the upstream model.
404pub(crate) fn target_model_or_source_model(state: &StreamTranslationState) -> String {
405    state
406        .target_model
407        .clone()
408        .or_else(|| state.model.clone())
409        .unwrap_or_else(|| "unknown".to_string())
410}
411
412// Returns the target/client ID when supplied, otherwise the upstream ID.
413pub(crate) fn target_message_id_or_source_message_id(
414    state: &StreamTranslationState,
415) -> Option<&str> {
416    state
417        .target_message_id
418        .as_deref()
419        .or(state.message_id.as_deref())
420}
421
422// Checks whether the current source format matches a built-in format.
423pub(crate) fn state_source_is(state: &StreamTranslationState, format: WireFormat) -> bool {
424    let format_id: FormatId = format.into();
425    match &state.source {
426        Some(source) => source == &format_id,
427        None => false,
428    }
429}
430
431// Reads a non-empty string field from an event object.
432pub(crate) fn string_field(object: &Map<String, Value>, key: &str) -> Option<String> {
433    object
434        .get(key)
435        .and_then(Value::as_str)
436        .filter(|value| !value.is_empty())
437        .map(ToOwned::to_owned)
438}
switchyard_translation/codecs/stream.rs

switchyard_translation/codecs/
stream.rs
