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    pub(crate) emitted_message_delta: bool,
46
47    pub(crate) next_content_index: usize,
48    pub(crate) text_block_index: Option<usize>,
49    pub(crate) text_block_started: bool,
50    pub(crate) emitted_content_block: bool,
51    pub(crate) tool_states: BTreeMap<usize, StreamToolState>,
52
53    pub(crate) response_created: bool,
54    pub(crate) response_text_started: bool,
55    pub(crate) response_text_output_index: Option<usize>,
56    pub(crate) response_text: String,
57    pub(crate) response_reasoning_started: bool,
58    pub(crate) response_reasoning_output_index: Option<usize>,
59    pub(crate) response_reasoning_text: String,
60    pub(crate) next_response_output_index: usize,
61
62    pub(crate) reasoning_block_index: Option<usize>,
63    pub(crate) reasoning_block_started: bool,
64}
65
66// Tracks an in-progress streamed tool call across provider-specific deltas.
67#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
68pub(crate) struct StreamToolState {
69    pub(crate) id: Option<String>,
70    pub(crate) name: Option<String>,
71    pub(crate) arguments: String,
72    pub(crate) pending_arguments: String,
73    pub(crate) started: bool,
74    pub(crate) content_index: Option<usize>,
75    pub(crate) response_output_index: Option<usize>,
76    pub(crate) response_item_id: Option<String>,
77}
78
79impl StreamTranslationState {
80    /// Creates stream state with source and target formats already attached.
81    pub fn new(source: impl Into<FormatId>, target: impl Into<FormatId>) -> Self {
82        Self {
83            source: Some(source.into()),
84            target: Some(target.into()),
85            ..Self::default()
86        }
87    }
88}
89
90/// Registry-backed streaming translator.
91#[derive(Default)]
92pub struct StreamTranslationEngine {
93    engine: TranslationEngine,
94}
95
96/// Codec contract for one provider streaming event format.
97pub trait StreamCodec: Send + Sync {
98    /// Returns the stream format handled by this codec.
99    fn format(&self) -> FormatId;
100
101    /// Decodes one provider event into zero or more neutral events.
102    fn decode_event(
103        &self,
104        state: &mut StreamTranslationState,
105        event: &Value,
106    ) -> Vec<LlmResponseChunk>;
107
108    /// Encodes one neutral event into zero or more provider events.
109    fn encode_event(
110        &self,
111        state: &mut StreamTranslationState,
112        event: LlmResponseChunk,
113    ) -> Vec<Value>;
114
115    /// Advances encoder state after an exact same-format event replay.
116    ///
117    /// Exact replay returns the preserved provider JSON instead of the JSON emitted by
118    /// [`Self::encode_event`]. The encoder must nevertheless observe the normalized chunks so
119    /// [`Self::finish`] can close an incomplete stream without duplicating an already replayed
120    /// terminal event. Codecs whose terminal state cannot be inferred from `MessageStop` alone
121    /// may override this hook and inspect `raw`.
122    fn observe_replayed_event(
123        &self,
124        state: &mut StreamTranslationState,
125        _raw: &Value,
126        normalized: Vec<LlmResponseChunk>,
127    ) {
128        let replayed_terminal = normalized
129            .iter()
130            .any(|chunk| matches!(chunk, LlmResponseChunk::MessageStop { .. }));
131        for chunk in normalized {
132            drop(self.encode_event(state, chunk));
133        }
134        if replayed_terminal {
135            state.finished = true;
136        }
137    }
138
139    /// Emits any terminal provider events needed after the source stream ends.
140    ///
141    /// This is intentionally required on every codec. Some target formats
142    /// need explicit terminal events after the source closes (for example,
143    /// Anthropic ``message_delta``/``message_stop`` or Responses
144    /// ``response.completed``). Formats that have no source-close work should
145    /// return an empty vector explicitly so the no-op behavior is a conscious
146    /// codec-level choice.
147    fn finish(&self, state: &mut StreamTranslationState) -> Vec<Value>;
148}
149
150/// Registry mapping stream wire formats to stream codecs.
151#[derive(Default)]
152pub struct StreamCodecRegistry {
153    codecs: BTreeMap<FormatId, Arc<dyn StreamCodec>>,
154}
155
156impl StreamCodecRegistry {
157    /// Creates an empty stream codec registry.
158    pub fn new() -> Self {
159        Self::default()
160    }
161
162    /// Creates a registry populated with built-in stream codecs.
163    pub fn with_builtins() -> Self {
164        let mut registry = Self::new();
165        registry.register(OpenAiChatStreamCodec);
166        registry.register(AnthropicMessagesStreamCodec);
167        registry.register(OpenAiResponsesStreamCodec);
168        registry
169    }
170
171    /// Registers or replaces a stream codec for its declared format.
172    pub fn register(&mut self, codec: impl StreamCodec + 'static) {
173        self.codecs.insert(codec.format(), Arc::new(codec));
174    }
175
176    /// Looks up a stream codec by format identifier.
177    pub fn codec(&self, format: impl Into<FormatId>) -> Result<Arc<dyn StreamCodec>> {
178        let format = format.into();
179        self.codecs.get(&format).cloned().ok_or_else(|| {
180            TranslationError::Other(format!("no stream codec registered for {format}"))
181        })
182    }
183}
184
185impl StreamTranslationEngine {
186    /// Creates a streaming engine from an explicit codec registry.
187    pub fn new(registry: StreamCodecRegistry) -> Self {
188        Self {
189            engine: TranslationEngine::with_registries(FormatRegistry::with_builtins(), registry),
190        }
191    }
192
193    /// Translates one source provider event into target provider events.
194    pub fn translate_event(
195        &self,
196        state: &mut StreamTranslationState,
197        source: impl Into<FormatId>,
198        target: impl Into<FormatId>,
199        event: &Value,
200    ) -> Result<Vec<Value>> {
201        self.engine.translate_event(state, source, target, event)
202    }
203
204    /// Finishes target-provider stream emission after the source stream closes.
205    pub fn finish(
206        &self,
207        state: &mut StreamTranslationState,
208        target: impl Into<FormatId>,
209    ) -> Result<Vec<Value>> {
210        self.engine.finish_stream(state, target)
211    }
212
213    /// Convenience helper using built-in codecs and error events instead of `Result`.
214    pub fn translate_event_with_builtins(
215        state: &mut StreamTranslationState,
216        source: impl Into<FormatId>,
217        target: impl Into<FormatId>,
218        event: &Value,
219    ) -> Vec<Value> {
220        Self::default()
221            .translate_event(state, source, target, event)
222            .unwrap_or_else(|error| vec![json!({"error": {"message": error.to_string()}})])
223    }
224}
225
226/// Decodes one provider stream event with the built-in codec registry.
227pub fn decode_stream_event(
228    state: &mut StreamTranslationState,
229    source: impl Into<FormatId>,
230    event: &Value,
231) -> Vec<LlmResponseChunk> {
232    let source = source.into();
233    StreamCodecRegistry::with_builtins()
234        .codec(source)
235        .map(|codec| codec.decode_event(state, event))
236        .unwrap_or_else(|error| {
237            // A missing codec is a translation-side failure, not something the
238            // upstream sent, so it decodes to `DecodeError` rather than `StreamError`.
239            vec![LlmResponseChunk::DecodeError {
240                message: error.to_string(),
241            }]
242        })
243}
244
245pub(crate) fn encode_response_stream_event(
246    state: &mut StreamTranslationState,
247    target_codec: &dyn StreamCodec,
248    target: &FormatId,
249    event: crate::LlmResponseStreamEvent,
250) -> Vec<Value> {
251    if state.errored {
252        return Vec::new();
253    }
254    let (preservation, normalized) = event.into_parts();
255    if let Some(preservation) = preservation {
256        let (source, raw) = preservation.into_parts();
257        if &source == target {
258            // Exact replay bypasses the target encoder's emitted JSON, but the encoder must
259            // still observe every normalized chunk. Otherwise `finish` starts from empty state:
260            // a clean EOF after a nonterminal provider event can omit or synthesize malformed
261            // terminal events, while a replayed terminal can be emitted twice.
262            target_codec.observe_replayed_event(state, &raw, normalized);
263            return vec![raw];
264        }
265    }
266
267    normalized
268        .into_iter()
269        .flat_map(|chunk| target_codec.encode_event(state, chunk))
270        .collect()
271}
272
273/// Encodes one neutral stream event with the built-in codec registry.
274pub fn encode_stream_event(
275    state: &mut StreamTranslationState,
276    target: impl Into<FormatId>,
277    event: LlmResponseChunk,
278) -> Vec<Value> {
279    StreamCodecRegistry::with_builtins()
280        .codec(target)
281        .map(|codec| codec.encode_event(state, event))
282        .unwrap_or_else(|error| vec![json!({"error": {"message": error.to_string()}})])
283}
284
285// Records source-provider identity carried by decoded stream events.
286pub(crate) fn record_source_identity(
287    state: &mut StreamTranslationState,
288    id: Option<String>,
289    model: Option<String>,
290) {
291    if id.is_some() {
292        state.message_id = id;
293    }
294    if model.is_some() {
295        state.model = model;
296    }
297}
298
299// Returns the served model when supplied, otherwise the upstream model.
300pub(crate) fn target_model_or_source_model(state: &StreamTranslationState) -> String {
301    state
302        .target_model
303        .clone()
304        .or_else(|| state.model.clone())
305        .unwrap_or_else(|| "unknown".to_string())
306}
307
308// Returns the target/client ID when supplied, otherwise the upstream ID.
309pub(crate) fn target_message_id_or_source_message_id(
310    state: &StreamTranslationState,
311) -> Option<&str> {
312    state
313        .target_message_id
314        .as_deref()
315        .or(state.message_id.as_deref())
316}
317
318// Checks whether the current source format matches a built-in format.
319pub(crate) fn state_source_is(state: &StreamTranslationState, format: WireFormat) -> bool {
320    let format_id: FormatId = format.into();
321    match &state.source {
322        Some(source) => source == &format_id,
323        None => false,
324    }
325}
326
327// Reads a non-empty string field from an event object.
328pub(crate) fn string_field(object: &Map<String, Value>, key: &str) -> Option<String> {
329    object
330        .get(key)
331        .and_then(Value::as_str)
332        .filter(|value| !value.is_empty())
333        .map(ToOwned::to_owned)
334}
switchyard_translation/codecs/stream.rs

switchyard_translation/codecs/
stream.rs
