dwara_core/ai/stream.rs
1//! Streaming translation (DW-077): provider SSE in, OpenAI SSE out,
2//! frame by frame.
3//!
4//! The [`StreamTranslator`] is the pure half of the streaming path:
5//! feed it the provider's response BYTES as they arrive and it returns
6//! the client-facing SSE TEXT to forward for exactly those bytes —
7//! nothing is held: a complete provider frame becomes complete client
8//! frames in the same poll, and a partial frame stays in the
9//! [`SseDecoder`] buffer (bounded by the provider's own frame size)
10//! until its terminator arrives.
11//!
12//! The dataplane's `AiStreamBody` (see `dataplane/ai_proxy.rs`) drives
13//! this translator over the upstream body and owns the metrics hooks.
14//!
15//! # Client-facing stream shape
16//!
17//! The client is OpenAI-shaped regardless of provider dialect:
18//!
19//! - each provider delta becomes one `data: {chat.completion.chunk}`
20//! frame (via [`openai_compat::stream_event_to_openai_chunk`]),
21//! - provider `Usage` events are ACCUMULATED (not forwarded inline) —
22//! they merge across the stream (Anthropic reports input tokens at
23//! `message_start` and output tokens at `message_delta`) —
24//! - at stream end the accumulated usage (if any) is emitted as ONE
25//! terminal `choices: []` usage chunk (the `stream_options.
26//! include_usage` shape), and the gateway ALWAYS writes its own
27//! `data: [DONE]` terminator (the provider's own terminator is
28//! swallowed, even OpenAI's, so the terminal ordering is uniform:
29//! deltas, usage, DONE).
30//!
31//! # Token counting (the locked decision)
32//!
33//! Token counts are PROVIDER-REPORTED ONLY. The translator accumulates
34//! the usage events the provider emits mid-stream; there is no local
35//! estimation, and chunk counts / text lengths are never presented as
36//! token estimates.
37//!
38//! # Mid-stream aborts
39//!
40//! [`StreamTranslator::abort_tail`] produces the terminal frames for a
41//! provider stream that died after frames were already forwarded: an
42//! OpenAI-shaped error chunk then `data: [DONE]` — the stream ends
43//! cleanly for the client instead of a connection reset, and the
44//! already-forwarded content stands.
45
46use crate::ai::adapter::ProviderAdapter;
47use crate::ai::openai_compat;
48use crate::ai::sse::{SseDecoder, SseFrame};
49use crate::ai::types::{StreamEvent, Usage};
50use serde_json::Value;
51
52/// One translated client frame (raw SSE text, terminator included).
53pub type ClientFrame = String;
54
55/// Pure streaming translator state (DW-077).
56pub struct StreamTranslator {
57 decoder: SseDecoder,
58 /// Provider-reported usage accumulated across the stream.
59 usage: Usage,
60 /// The provider signalled end-of-stream (its terminator event, its
61 /// `[DONE]` sentinel, or the body's clean end).
62 ended: bool,
63 /// The terminal frames (usage chunk + gateway [DONE]) were
64 /// emitted. Distinct from `ended`: the provider's own terminator
65 /// sets `ended` mid-stream, and the terminal frames must STILL be
66 /// emitted afterwards — exactly once.
67 tail_emitted: bool,
68 /// Identity stamped on every client chunk.
69 id: String,
70 model_alias: String,
71 created: u64,
72}
73
74impl StreamTranslator {
75 /// New translator stamping chunks with `id` (the response id),
76 /// `model_alias` (what the client asked for — never the provider
77 /// model), and `created` (unix seconds, fixed at stream start).
78 pub fn new(id: String, model_alias: String, created: u64) -> Self {
79 StreamTranslator {
80 decoder: SseDecoder::new(),
81 usage: Usage::default(),
82 ended: false,
83 tail_emitted: false,
84 id,
85 model_alias,
86 created,
87 }
88 }
89
90 /// Feed provider bytes; returns the client frames for every
91 /// COMPLETE provider frame these bytes complete (may be empty —
92 /// bytes that only continue a partial frame). Also reports how
93 /// many delta chunks were forwarded and whether one of them was
94 /// the FIRST forwarded chunk of the stream (the first-token
95 /// signal).
96 pub fn feed(
97 &mut self,
98 bytes: &[u8],
99 adapter: &dyn ProviderAdapter,
100 ) -> (Vec<ClientFrame>, usize, bool) {
101 let frames = self.decoder.push(bytes);
102 self.translate(frames, adapter)
103 }
104
105 /// Flush a final provider frame not terminated by a blank line
106 /// (tolerant providers exist; well-formed streams yield nothing
107 /// here — the body's clean end drives [`Self::finish`]).
108 pub fn flush_partial(
109 &mut self,
110 adapter: &dyn ProviderAdapter,
111 ) -> (Vec<ClientFrame>, usize, bool) {
112 let frames = self.decoder.finish();
113 self.translate(frames, adapter)
114 }
115
116 fn translate(
117 &mut self,
118 frames: Vec<SseFrame>,
119 adapter: &dyn ProviderAdapter,
120 ) -> (Vec<ClientFrame>, usize, bool) {
121 let mut out = Vec::new();
122 let mut chunk_count = 0usize;
123 let mut first = false;
124 for frame in frames {
125 // The provider's own end-of-stream sentinel (OpenAI
126 // `[DONE]`): swallowed — the gateway writes its own
127 // terminator in `finish`.
128 if Some(frame.data.as_str()) == adapter.stream_done_sentinel() {
129 self.ended = true;
130 continue;
131 }
132 let Ok(data) = serde_json::from_str::<Value>(&frame.data) else {
133 // A non-JSON data line is a provider bug; skip the
134 // frame rather than kill the stream.
135 continue;
136 };
137 let Ok(events) = adapter.parse_stream_event(&data) else {
138 // Same posture as a non-JSON line: one bad frame must
139 // not take down an otherwise healthy stream.
140 continue;
141 };
142 for event in events {
143 match event {
144 StreamEvent::Delta(delta) => {
145 let chunk = openai_compat::stream_event_to_openai_chunk(
146 &StreamEvent::Delta(delta),
147 &self.id,
148 &self.model_alias,
149 self.created,
150 );
151 out.push(sse_data(&chunk));
152 chunk_count += 1;
153 first = true;
154 }
155 StreamEvent::Usage(u) => {
156 // Accumulated, not forwarded: one terminal
157 // usage chunk at stream end carries the merged
158 // provider-reported totals.
159 self.usage.merge(u);
160 }
161 StreamEvent::Done => {
162 self.ended = true;
163 }
164 }
165 }
166 }
167 // `first` is only true for the first TRANSLATED batch that
168 // contained a delta; the caller treats it as the first-token
169 // signal when it has not seen one yet.
170 (out, chunk_count, first)
171 }
172
173 /// The terminal frames for a CLEAN stream end: the accumulated
174 /// usage chunk (when the provider reported any usage) followed by
175 /// the gateway's own `data: [DONE]`. Returns them once; later
176 /// calls return empty (idempotent tail).
177 pub fn finish(&mut self) -> Vec<ClientFrame> {
178 if self.tail_emitted {
179 return Vec::new();
180 }
181 self.tail_emitted = true;
182 self.ended = true;
183 let mut out = Vec::new();
184 if self.usage.prompt_tokens.is_some() || self.usage.completion_tokens.is_some() {
185 let usage = openai_compat::stream_event_to_openai_chunk(
186 &StreamEvent::Usage(self.usage),
187 &self.id,
188 &self.model_alias,
189 self.created,
190 );
191 out.push(sse_data(&usage));
192 }
193 out.push("data: [DONE]\n\n".to_string());
194 out
195 }
196
197 /// The terminal frames for a MID-STREAM ABORT (the provider body
198 /// errored after frames were forwarded): an OpenAI-shaped error
199 /// chunk then the terminator. Already-forwarded content stands.
200 pub fn abort_tail(&mut self, message: &str) -> Vec<ClientFrame> {
201 self.tail_emitted = true;
202 self.ended = true;
203 vec![
204 sse_data(&openai_compat::error_body(
205 message,
206 "api_error",
207 Some("provider_stream_aborted"),
208 &self.id,
209 )),
210 "data: [DONE]\n\n".to_string(),
211 ]
212 }
213
214 /// Whether the provider signalled end-of-stream (its terminator or
215 /// the body's clean end already consumed).
216 pub fn is_ended(&self) -> bool {
217 self.ended
218 }
219
220 /// The accumulated provider-reported usage so far.
221 pub fn usage(&self) -> Usage {
222 self.usage
223 }
224
225 /// The response id stamped on every client chunk (AI-07: the
226 /// `gen_ai.response.id` value for streaming, where the gateway
227 /// synthesizes the id the provider does not send as one frame).
228 pub fn id(&self) -> &str {
229 &self.id
230 }
231}
232
233/// Serialize one JSON value as an SSE `data:` frame (event streams to
234/// OpenAI clients carry only data lines).
235fn sse_data(v: &Value) -> String {
236 format!("data: {}\n\n", v)
237}
238
239#[cfg(test)]
240mod tests {
241 use super::*;
242 use crate::ai::adapter::adapter_for;
243 use crate::config::ai::AiProviderKind;
244
245 // White-box: the translator is the private half of the streaming
246 // pipeline; the gateway-level behavior is covered end to end by
247 // tests/ai_streaming.rs. These pin the frame grammar cheaply.
248
249 #[test]
250 fn openai_deltas_translate_and_terminator_is_swallowed() {
251 let mut t = StreamTranslator::new("chatcmpl-x".into(), "alias".into(), 7);
252 let adapter = adapter_for(AiProviderKind::Openai);
253 let (frames, n, first) = t.feed(
254 concat!(
255 "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"he\"}}]}\n\n",
256 "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"llo\"}}]}\n\n"
257 )
258 .as_bytes(),
259 adapter,
260 );
261 assert_eq!(n, 2);
262 assert!(first);
263 assert_eq!(frames.len(), 2);
264 assert!(frames[0].starts_with("data: {"));
265 assert!(frames[0].contains("chat.completion.chunk"));
266 assert!(frames[1].contains("llo"));
267 // The provider's [DONE] is swallowed; OUR finish writes it.
268 let (_, n2, _) = t.feed(b"data: [DONE]\n\n", adapter);
269 assert_eq!(n2, 0);
270 assert!(t.is_ended());
271 }
272
273 #[test]
274 fn usage_accumulates_and_finish_emits_one_terminal_chunk() {
275 let mut t = StreamTranslator::new("c".into(), "m".into(), 1);
276 let adapter = adapter_for(AiProviderKind::Anthropic);
277 // input tokens at message_start, output tokens at message_delta
278 t.feed(
279 b"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":100}}}\n\n",
280 adapter,
281 );
282 t.feed(
283 b"event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":50}}\n\n",
284 adapter,
285 );
286 let tail = t.finish();
287 assert_eq!(tail.len(), 2);
288 assert!(tail[0].contains("\"prompt_tokens\":100"));
289 assert!(tail[0].contains("\"completion_tokens\":50"));
290 assert!(tail[0].contains("\"choices\":[]"));
291 assert_eq!(tail[1], "data: [DONE]\n\n");
292 // Idempotent: a second finish yields nothing.
293 assert!(t.finish().is_empty());
294 }
295
296 #[test]
297 fn abort_tail_shapes_an_error_chunk_then_done() {
298 let mut t = StreamTranslator::new("c".into(), "m".into(), 1);
299 let tail = t.abort_tail("the model provider closed the stream");
300 assert_eq!(tail.len(), 2);
301 assert!(tail[0].contains("provider_stream_aborted"));
302 assert!(tail[0].contains("the model provider closed the stream"));
303 assert_eq!(tail[1], "data: [DONE]\n\n");
304 }
305
306 #[test]
307 fn bad_frames_are_skipped_not_fatal() {
308 let mut t = StreamTranslator::new("c".into(), "m".into(), 1);
309 let adapter = adapter_for(AiProviderKind::Openai);
310 let (frames, n, _) = t.feed(b"data: not-json\n\n", adapter);
311 assert_eq!(n, 0);
312 assert!(frames.is_empty());
313 // The stream continues after the bad frame.
314 let (frames2, n2, _) = t.feed(
315 b"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"}}]}\n\n",
316 adapter,
317 );
318 assert_eq!(n2, 1);
319 assert_eq!(frames2.len(), 1);
320 }
321}