harness_gateway_client/wire/
stream.rs

1//! SSE reassembly for the always-streaming completion protocol.
2//!
3//! [`SseScanner`] splits the raw byte stream into `data:` payloads, and
4//! [`StreamAccumulator`] folds those payloads back into the buffered
5//! chat-completion body shape, then [`finishes`](StreamAccumulator::finish)
6//! it into a [`Completion`]. The strict turn rules stay in [`super::parse`]:
7//! the accumulator only reassembles, so streamed and buffered turns are
8//! judged by one rule set.
9//!
10//! No HTTP happens here. The transport that reads the bytes off the wire
11//! hands them to the read loop, which feeds the scanner, hands payloads to
12//! the accumulator, and takes the completion from `finish`.
13//!
14//! The progress subscription in the model vocabulary deliberately has its
15//! own SSE decoder, and neither can substitute for the other: that one
16//! decodes blank-line-terminated event blocks into typed progress items and
17//! stays lossy (an undecodable block is one `Err` item in a telemetry
18//! stream), while this one hands raw `data:` payloads to a transport loop
19//! that meters bytes and timing and hard-fails on the first malformed
20//! chunk, because a completion's product must be whole.
21
22use std::collections::BTreeMap;
23
24use promptforge::metrics::{CallMetrics, ClientTiming};
25use promptforge::model::{Completion, CompletionError, RawExchange};
26use serde_json::{Map, Value};
27
28use super::classify::classify_stream_error;
29use super::delta::StreamDelta;
30use super::parse::{normalize, response_metadata};
31use crate::failure::malformed;
32
33/// Splits a raw SSE byte stream into `data:` payloads.
34///
35/// Blank lines, `:` comments, and non-`data:` fields (`event:`, `id:`,
36/// `retry:`) are skipped; the caller sees only payload text.
37///
38/// Crate-private: a transport reaches it only through
39/// [`read_completion_stream`](crate::read_completion_stream).
40#[derive(Debug, Default)]
41pub(super) struct SseScanner {
42    buffer: Vec<u8>,
43    /// How much of `buffer` is already known to hold no `\n`.
44    scanned: usize,
45}
46
47impl SseScanner {
48    /// A scanner with an empty buffer.
49    #[must_use]
50    pub(super) fn new() -> SseScanner {
51        SseScanner {
52            buffer: Vec::new(),
53            scanned: 0,
54        }
55    }
56
57    /// Buffers freshly received bytes for line extraction.
58    pub(super) fn extend(&mut self, bytes: &[u8]) {
59        self.buffer.extend_from_slice(bytes);
60    }
61
62    /// Returns the next complete `data:` payload, or `None` until one is
63    /// fully buffered.
64    pub(super) fn next_data(&mut self) -> Option<String> {
65        loop {
66            let Some(offset) = self.buffer[self.scanned..]
67                .iter()
68                .position(|byte| *byte == b'\n')
69            else {
70                self.scanned = self.buffer.len();
71                return None;
72            };
73            let end = self.scanned + offset;
74            self.scanned = 0;
75            let line: Vec<u8> = self.buffer.drain(..=end).collect();
76            let line = String::from_utf8_lossy(&line);
77            let line = line.trim_end_matches(['\r', '\n']);
78            if line.is_empty() || line.starts_with(':') {
79                continue;
80            }
81            let Some(data) = line.strip_prefix("data:") else {
82                continue;
83            };
84            return Some(data.trim_start().to_owned());
85        }
86    }
87}
88
89/// The outcome of applying one `data:` payload.
90#[derive(Debug, Clone, Copy, PartialEq, Eq)]
91pub(super) enum Applied {
92    /// The payload advanced the accumulation; `delta` is true when it
93    /// included answer text, reasoning, or a tool-call fragment (the
94    /// TTFT/ITL clock ticks on those, never on role or summary chunks).
95    Chunk {
96        /// Whether the chunk included generated content.
97        delta: bool,
98    },
99    /// The payload was the terminal `[DONE]` sentinel.
100    Done,
101}
102
103/// One tool call assembled from streamed fragments, keyed by the fragment
104/// `index`. `id`, `name`, and `arguments` each grow by string concatenation
105/// as fragments arrive, per the `OpenAI` streaming contract.
106#[derive(Debug, Default)]
107struct ToolCallParts {
108    id: String,
109    name: String,
110    arguments: String,
111}
112
113/// Accumulates streamed chunks into the buffered chat-completion shape.
114///
115/// Only the first choice (`index == 0`) is accumulated, mirroring the
116/// buffered normalizer, which reads `choices[0]` alone. Metadata sections
117/// (`usage`, llama.cpp `timings`, vLLM `metrics`) are kept verbatim from
118/// whichever chunk held them last, including the empty-choices summary
119/// chunk `stream_options.include_usage` appends, and are handed to the
120/// lenient metadata parser unjudged.
121///
122/// Crate-private: a transport reaches it only through
123/// [`read_completion_stream`](crate::read_completion_stream).
124#[derive(Debug, Default)]
125pub(super) struct StreamAccumulator {
126    /// Answer text; `None` until the first `content` fragment arrives.
127    content: Option<String>,
128    /// Reasoning side-channel text; `None` until the first fragment.
129    reasoning: Option<String>,
130    tool_calls: BTreeMap<u64, ToolCallParts>,
131    finish_reason: Option<String>,
132    model: Option<String>,
133    /// Raw top-level metadata sections, latest occurrence wins.
134    sections: Map<String, Value>,
135}
136
137impl StreamAccumulator {
138    /// An empty accumulator.
139    #[must_use]
140    pub(super) fn new() -> StreamAccumulator {
141        StreamAccumulator::default()
142    }
143
144    /// Applies one `data:` payload, invoking `on_delta` for each text or
145    /// reasoning fragment it contains.
146    ///
147    /// # Errors
148    /// Returns a `MalformedResponse`-kind [`CompletionError`] when the
149    /// payload is not valid JSON or a recognized field has the wrong
150    /// shape, and the [`classify_stream_error`] result when the payload is
151    /// a mid-stream error envelope (`Transport` unless its text names a
152    /// known cause).
153    pub(super) fn apply(
154        &mut self,
155        data: &str,
156        on_delta: &impl Fn(StreamDelta),
157    ) -> Result<Applied, CompletionError> {
158        if data == "[DONE]" {
159            return Ok(Applied::Done);
160        }
161        let chunk: Value = serde_json::from_str(data)
162            .map_err(|error| malformed("stream chunk was not valid JSON").with_source(error))?;
163        // A mid-stream `error` envelope is how the gateway (and llama.cpp)
164        // report a failure after the 200 has already been sent: the
165        // completion died in flight, so it is a transport failure unless
166        // the bounded, control-escaped message names a known cause.
167        if let Some(envelope) = chunk.get("error").filter(|error| !error.is_null()) {
168            let message = envelope
169                .get("message")
170                .and_then(Value::as_str)
171                .unwrap_or("stream error envelope omitted its message");
172            return Err(classify_stream_error(&escape_controls(message, 2000)));
173        }
174        if let Some(Value::String(model)) = chunk.get("model")
175            && !model.is_empty()
176        {
177            self.model = Some(model.clone());
178        }
179        for key in ["usage", "timings", "metrics"] {
180            if let Some(section) = chunk.get(key)
181                && !section.is_null()
182            {
183                self.sections.insert(key.to_owned(), section.clone());
184            }
185        }
186        // Absent or empty `choices` is the summary-chunk shape
187        // (`stream_options.include_usage`): metadata only, nothing to index.
188        let choices = match chunk.get("choices") {
189            None | Some(Value::Null) => return Ok(Applied::Chunk { delta: false }),
190            Some(Value::Array(choices)) => choices,
191            Some(_) => {
192                return Err(malformed(
193                    "stream chunk `choices` was present but not an array",
194                ));
195            }
196        };
197        let mut held_delta = false;
198        for choice in choices {
199            if self.apply_choice(choice, on_delta)? {
200                held_delta = true;
201            }
202        }
203        Ok(Applied::Chunk { delta: held_delta })
204    }
205
206    /// Applies one streamed choice, returning whether it held content.
207    fn apply_choice(
208        &mut self,
209        choice: &Value,
210        on_delta: &impl Fn(StreamDelta),
211    ) -> Result<bool, CompletionError> {
212        let Some(index) = choice.get("index").and_then(Value::as_u64) else {
213            return Err(malformed("stream choice had no integer index"));
214        };
215        // Mirror the buffered normalizer: the first choice is the turn.
216        if index != 0 {
217            return Ok(false);
218        }
219        match choice.get("finish_reason") {
220            None | Some(Value::Null) => {}
221            Some(Value::String(reason)) => self.finish_reason = Some(reason.clone()),
222            Some(_) => {
223                return Err(malformed(
224                    "stream choice `finish_reason` was present but not a string",
225                ));
226            }
227        }
228        let delta = match choice.get("delta") {
229            // A finish-only chunk may omit the delta entirely.
230            None | Some(Value::Null) => return Ok(false),
231            Some(delta @ Value::Object(_)) => delta,
232            Some(_) => {
233                return Err(malformed(
234                    "stream choice `delta` was present but not an object",
235                ));
236            }
237        };
238        let mut held = false;
239        if let Some(text) = append_string_fragment(delta, "content", &mut self.content, "content")?
240            && !text.is_empty()
241        {
242            held = true;
243            on_delta(StreamDelta::Text(text));
244        }
245        for key in ["reasoning_content", "reasoning", "thinking"] {
246            if let Some(text) = append_string_fragment(delta, key, &mut self.reasoning, key)?
247                && !text.is_empty()
248            {
249                held = true;
250                on_delta(StreamDelta::Reasoning(text));
251            }
252        }
253        match delta.get("tool_calls") {
254            None | Some(Value::Null) => {}
255            Some(Value::Array(fragments)) => {
256                for fragment in fragments {
257                    self.apply_tool_fragment(fragment)?;
258                }
259                if !fragments.is_empty() {
260                    held = true;
261                }
262            }
263            Some(_) => {
264                return Err(malformed(
265                    "stream delta `tool_calls` was present but not an array",
266                ));
267            }
268        }
269        Ok(held)
270    }
271
272    /// Merges one tool-call fragment into its index-keyed buffer.
273    fn apply_tool_fragment(&mut self, fragment: &Value) -> Result<(), CompletionError> {
274        let Some(index) = fragment.get("index").and_then(Value::as_u64) else {
275            return Err(malformed("stream tool-call fragment had no integer index"));
276        };
277        let parts = self.tool_calls.entry(index).or_default();
278        match fragment.get("id") {
279            None | Some(Value::Null) => {}
280            Some(Value::String(id)) => parts.id.push_str(id),
281            Some(_) => {
282                return Err(malformed("stream tool-call fragment `id` was not a string"));
283            }
284        }
285        let function = match fragment.get("function") {
286            None | Some(Value::Null) => return Ok(()),
287            Some(function @ Value::Object(_)) => function,
288            Some(_) => {
289                return Err(malformed(
290                    "stream tool-call fragment `function` was not an object",
291                ));
292            }
293        };
294        for (key, slot) in [
295            ("name", &mut parts.name),
296            ("arguments", &mut parts.arguments),
297        ] {
298            match function.get(key) {
299                None | Some(Value::Null) => {}
300                Some(Value::String(piece)) => slot.push_str(piece),
301                Some(_) => {
302                    return Err(malformed(format!(
303                        "stream tool-call fragment `{key}` was not a string"
304                    )));
305                }
306            }
307        }
308        Ok(())
309    }
310
311    /// Finishes the accumulation into the [`Completion`] the turn produced:
312    /// the truncation rule, the strict turn normalizer, and the lenient
313    /// metadata parser, in that order. `request_body` is the body the
314    /// transport sent; the completion is labeled with the model it names,
315    /// in place of the name the response gave, and with the reassembled
316    /// response it becomes the completion's [`RawExchange`] for the debug
317    /// capture. `client_timing`
318    /// is what the transport measured on its own clock; it joins the
319    /// backend's sections in the completion's [`CallMetrics`], which is
320    /// absent when nothing was measured.
321    ///
322    /// # Errors
323    /// Returns a `MalformedResponse`-kind [`CompletionError`] when a
324    /// tool-call batch was cut short by a `length` or `content_filter`
325    /// finish (partial arguments must not execute), the normalizer's own
326    /// errors (`EmptyReply` for a turn with neither non-empty tool calls nor
327    /// non-empty text), and the validating constructor's (a
328    /// `MalformedResponse` for two calls sharing an id).
329    pub(super) fn finish(
330        self,
331        request_body: Value,
332        client_timing: Option<ClientTiming>,
333    ) -> Result<Completion, CompletionError> {
334        // The truncation rule runs before normalization: a tool-call batch
335        // cut short by `length` or `content_filter` may hold partial JSON
336        // arguments, and partial arguments must not execute.
337        if !self.tool_calls.is_empty()
338            && matches!(
339                self.finish_reason.as_deref(),
340                Some("length" | "content_filter")
341            )
342        {
343            let reason = self.finish_reason.unwrap_or_default();
344            return Err(malformed(format!(
345                "tool-call batch truncated by finish_reason {reason:?}: \
346                 partial arguments must not execute"
347            )));
348        }
349        let response_body = self.into_body();
350        let turn = normalize(&response_body)?;
351        let metadata = response_metadata(&response_body);
352        let metrics = CallMetrics {
353            client: client_timing,
354            ..metadata.metrics
355        };
356        let measured = metrics.usage.is_some()
357            || metrics.llama.is_some()
358            || metrics.vllm.is_some()
359            || metrics.client.is_some();
360        let mut completion = Completion::from_result(turn.outcome, metadata.model)?
361            .with_metadata_diagnostics(metadata.diagnostics);
362        if let Some(Value::String(model)) = request_body.get("model") {
363            completion = completion.with_model(model.clone());
364        }
365        if let Some(reason) = turn.finish_reason {
366            completion = completion.with_finish_reason(reason);
367        }
368        if let Some(reasoning) = turn.reasoning_content {
369            completion = completion.with_reasoning_content(reasoning);
370        }
371        if measured {
372            completion = completion.with_metrics(metrics);
373        }
374        Ok(completion.with_raw(RawExchange::new(request_body, response_body)))
375    }
376
377    /// Reassembles the accumulation into the buffered chat-completion body
378    /// shape, ready for the strict turn normalizer and the lenient metadata
379    /// parser.
380    fn into_body(self) -> Value {
381        let mut message = Map::new();
382        message.insert("role".to_owned(), Value::String("assistant".to_owned()));
383        message.insert(
384            "content".to_owned(),
385            match self.content {
386                Some(text) => Value::String(text),
387                None => Value::Null,
388            },
389        );
390        if let Some(reasoning) = self.reasoning.filter(|text| !text.is_empty()) {
391            message.insert("reasoning_content".to_owned(), Value::String(reasoning));
392        }
393        if !self.tool_calls.is_empty() {
394            let calls: Vec<Value> = self
395                .tool_calls
396                .into_values()
397                .map(|parts| {
398                    serde_json::json!({
399                        "id": parts.id,
400                        "type": "function",
401                        "function": { "name": parts.name, "arguments": parts.arguments },
402                    })
403                })
404                .collect();
405            message.insert("tool_calls".to_owned(), Value::Array(calls));
406        }
407        let mut choice = Map::new();
408        choice.insert("index".to_owned(), Value::from(0));
409        choice.insert("message".to_owned(), Value::Object(message));
410        if let Some(reason) = self.finish_reason {
411            choice.insert("finish_reason".to_owned(), Value::String(reason));
412        }
413        let mut body = Map::new();
414        if let Some(model) = self.model {
415            body.insert("model".to_owned(), Value::String(model));
416        }
417        body.insert(
418            "choices".to_owned(),
419            Value::Array(vec![Value::Object(choice)]),
420        );
421        for (key, value) in self.sections {
422            body.insert(key, value);
423        }
424        Value::Object(body)
425    }
426}
427
428/// Appends a string fragment under `key` from `delta` into `slot`,
429/// returning the fragment when one was present.
430///
431/// Absent and JSON-null are no fragment; a present non-string is a
432/// malformed shape named after `label`.
433fn append_string_fragment(
434    delta: &Value,
435    key: &str,
436    slot: &mut Option<String>,
437    label: &str,
438) -> Result<Option<String>, CompletionError> {
439    match delta.get(key) {
440        None | Some(Value::Null) => Ok(None),
441        Some(Value::String(text)) => {
442            slot.get_or_insert_with(String::new).push_str(text);
443            Ok(Some(text.clone()))
444        }
445        Some(_) => Err(malformed(format!(
446            "stream delta `{label}` was present but not a string"
447        ))),
448    }
449}
450
451/// Truncates a backend response body and escapes its control characters so it
452/// is safe to show in a diagnostic.
453///
454/// The function keeps the first `max` characters of `body`. It replaces each
455/// control character among them, including newlines and carriage returns, with
456/// its escaped form, such as `\n` or `\u{1b}`. Escaping can make the result
457/// longer than `max` characters. An empty body returns the fixed marker
458/// `(empty body)`.
459///
460/// The escaping stops a backend body from forging log lines or slipping
461/// terminal control sequences into a diagnostic.
462///
463/// A transport passes the error body of a response with a failure status
464/// through this function before it hands the body to
465/// [`classify_http_failure`](crate::classify_http_failure). Every transport
466/// therefore truncates and escapes a backend body by the same rule.
467#[must_use]
468pub fn escape_controls(body: &str, max: usize) -> String {
469    if body.is_empty() {
470        return "(empty body)".to_owned();
471    }
472    let mut escaped = String::with_capacity(body.len());
473    for ch in body.chars().take(max) {
474        if ch.is_control() {
475            for part in ch.escape_default() {
476                escaped.push(part);
477            }
478        } else {
479            escaped.push(ch);
480        }
481    }
482    escaped
483}
484
485#[cfg(test)]
486#[path = "stream-tests.rs"]
487mod tests;
488
489#[cfg(test)]
490#[path = "stream-metrics-tests.rs"]
491mod metrics_tests;
harness_gateway_client/wire/stream.rs

harness_gateway_client/wire/
stream.rs
