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

switchyard_translation/codecs/
stream.rs
