Skip to main content

gglib_core/sse/
encoder.rs

1//! Encode typed [`LlmStreamEvent`] values into OpenAI-compatible SSE
2//! `chat.completion.chunk` `data:` frames.
3//!
4//! This is the inverse of [`super::parser::parse_sse_frame`] and is used by
5//! the proxy after the universal normalization layer has rewritten model-
6//! specific dialects (Qwen XML tool calls, bare `<think>` tags) into strict
7//! `OpenAI` events.  Re-emitting the canonical wire format ensures external
8//! clients (`OpenWebUI`, `OpenAI` SDKs, etc.) see only pristine `OpenAI` JSON
9//! regardless of which model is on the other end.
10//!
11//! # Frame envelope
12//!
13//! Every emitted chunk has this shape:
14//!
15//! ```json
16//! {
17//!   "id": "chatcmpl-…",
18//!   "object": "chat.completion.chunk",
19//!   "created": 1729000000,
20//!   "model": "qwen3-coder",
21//!   "choices": [{ "index": 0, "delta": { … }, "finish_reason": null }]
22//! }
23//! ```
24//!
25//! Stable values (`id`, `model`, `created`) are carried on [`SseEncoder`] so
26//! they are identical across every chunk of a single response.
27
28use serde_json::{Value, json};
29
30use crate::LlmStreamEvent;
31
32/// The SSE stream-terminator sentinel.
33///
34/// Must be sent by the caller exactly once, only after the entire event
35/// stream from [`SseEncoder::encode`] is truly exhausted — never bundled
36/// into an individual event's encoding, since [`LlmStreamEvent::Done`] is
37/// not guaranteed to be the last event (a trailing
38/// [`LlmStreamEvent::Usage`] can legitimately follow it) and nothing may
39/// be sent after `[DONE]` on the wire.
40pub const DONE_SENTINEL: &str = "data: [DONE]\n\n";
41
42/// Stateful encoder that produces OpenAI-shape SSE frames for one response.
43///
44/// The `id`, `model`, and `created` fields are stable across all frames the
45/// encoder produces, matching the `OpenAI` streaming contract.
46#[derive(Debug, Clone)]
47pub struct SseEncoder {
48    /// Stable response id, e.g. `"chatcmpl-…"`.
49    pub id: String,
50    /// Model name as advertised to the client (NOT the upstream alias).
51    pub model: String,
52    /// Unix epoch seconds when the response was created.
53    pub created: u64,
54}
55
56impl SseEncoder {
57    /// Construct a new encoder with the stable response metadata.
58    #[must_use]
59    pub fn new(id: impl Into<String>, model: impl Into<String>, created: u64) -> Self {
60        Self {
61            id: id.into(),
62            model: model.into(),
63            created,
64        }
65    }
66
67    /// Encode a single [`LlmStreamEvent`] into one or more SSE frames.
68    ///
69    /// Returns `None` when the event is not meant to appear on the wire (e.g.
70    /// [`LlmStreamEvent::NormalizationError`], which the proxy logs but never
71    /// forwards to clients).
72    ///
73    /// For [`LlmStreamEvent::Done`], the returned `String` is only the
74    /// terminating chunk (with `finish_reason` set) — it deliberately does
75    /// **not** include the trailing `data: [DONE]\n\n` sentinel
76    /// ([`DONE_SENTINEL`]).  `Done` is not guaranteed to be the last event on
77    /// the wire: a trailing [`LlmStreamEvent::Usage`] can legitimately arrive
78    /// afterward (see that variant's doc), and nothing may follow `[DONE]`
79    /// once it's sent.  Callers must append [`DONE_SENTINEL`] themselves,
80    /// exactly once, only after the entire event stream is truly exhausted.
81    #[must_use]
82    pub fn encode(&self, event: &LlmStreamEvent) -> Option<String> {
83        match event {
84            LlmStreamEvent::TextDelta { content } => Some(self.frame(&json!({
85                "index": 0,
86                "delta": { "content": content },
87                "finish_reason": Value::Null,
88            }))),
89            LlmStreamEvent::ReasoningDelta { content } => Some(self.frame(&json!({
90                "index": 0,
91                "delta": { "reasoning_content": content },
92                "finish_reason": Value::Null,
93            }))),
94            LlmStreamEvent::ToolCallDelta {
95                index,
96                id,
97                name,
98                arguments,
99            } => {
100                let mut tc = json!({ "index": index });
101                if let Some(id) = id {
102                    tc["id"] = json!(id);
103                    // OpenAI clients expect "type":"function" on the first
104                    // delta for a given index.
105                    tc["type"] = json!("function");
106                }
107                let mut function = json!({});
108                if let Some(name) = name {
109                    function["name"] = json!(name);
110                }
111                if let Some(arguments) = arguments {
112                    function["arguments"] = json!(arguments);
113                }
114                if function.as_object().is_some_and(|o| !o.is_empty()) {
115                    tc["function"] = function;
116                }
117                Some(self.frame(&json!({
118                    "index": 0,
119                    "delta": { "tool_calls": [tc] },
120                    "finish_reason": Value::Null,
121                })))
122            }
123            LlmStreamEvent::PromptProgress {
124                processed,
125                total,
126                cached,
127                time_ms,
128            } => {
129                // prompt_progress frames live at the top level (no `choices`).
130                let value = json!({
131                    "id": self.id,
132                    "object": "chat.completion.chunk",
133                    "created": self.created,
134                    "model": self.model,
135                    "prompt_progress": {
136                        "processed": processed,
137                        "total": total,
138                        "cache": cached,
139                        "time_ms": time_ms,
140                    },
141                });
142                Some(format!("data: {value}\n\n"))
143            }
144            LlmStreamEvent::Done { finish_reason } => Some(self.frame(&json!({
145                "index": 0,
146                "delta": {},
147                "finish_reason": finish_reason,
148            }))),
149            LlmStreamEvent::Usage {
150                prompt_tokens,
151                completion_tokens,
152                total_tokens,
153                cached_tokens,
154            } => Some(self.usage_frame(
155                *prompt_tokens,
156                *completion_tokens,
157                *total_tokens,
158                *cached_tokens,
159            )),
160            LlmStreamEvent::NormalizationError { .. } => None,
161            LlmStreamEvent::UpstreamError {
162                message,
163                error_type,
164                code,
165            } => Some(Self::upstream_error_frame(message, error_type, code)),
166        }
167    }
168
169    /// Encode a [`LlmStreamEvent::Usage`] event.
170    ///
171    /// Per the `OpenAI` `stream_options.include_usage` convention, the
172    /// usage-totals chunk carries an empty `choices` array (not omitted —
173    /// see [`crate::LlmStreamEvent::Usage`] doc) and a top-level `usage`
174    /// object.
175    fn usage_frame(
176        &self,
177        prompt_tokens: u32,
178        completion_tokens: u32,
179        total_tokens: u32,
180        cached_tokens: Option<u32>,
181    ) -> String {
182        let mut usage = json!({
183            "prompt_tokens": prompt_tokens,
184            "completion_tokens": completion_tokens,
185            "total_tokens": total_tokens,
186        });
187        // Re-emitted only when the upstream reported it, so the frame stays
188        // byte-identical to before for servers that don't. Clients such as the
189        // Copilot LLM Gateway extension surface this as `promptTokenDetails`.
190        if let Some(cached) = cached_tokens {
191            usage["prompt_tokens_details"] = json!({ "cached_tokens": cached });
192        }
193        let value = json!({
194            "id": self.id,
195            "object": "chat.completion.chunk",
196            "created": self.created,
197            "model": self.model,
198            "choices": [],
199            "usage": usage,
200        });
201        format!("data: {value}\n\n")
202    }
203
204    /// Encode a [`LlmStreamEvent::UpstreamError`] event.
205    ///
206    /// Deliberately bare — no `id`/`object`/`created`/`model` envelope and,
207    /// crucially, no `choices` key at all (unlike every other frame this
208    /// encoder produces). Clients such as the GitHub Copilot LLM Gateway
209    /// extension detect this exact shape
210    /// (`'error' in obj && !('choices' in obj)`) to recognise an inline
211    /// mid-stream failure; wrapping it in the usual envelope or adding an
212    /// empty `choices: []` would hide it as an ordinary chunk instead.
213    ///
214    /// Does **not** append [`DONE_SENTINEL`] — see [`Self::encode`] doc; the
215    /// caller appends it exactly once after the stream is truly exhausted.
216    fn upstream_error_frame(message: &str, error_type: &str, code: &str) -> String {
217        let error_obj = json!({
218            "error": {
219                "message": message,
220                "type": error_type,
221                "code": code,
222            }
223        });
224        format!("data: {error_obj}\n\n")
225    }
226
227    /// Wrap a `choice` value in the standard chunk envelope and SSE framing.
228    fn frame(&self, choice: &Value) -> String {
229        let value = json!({
230            "id": self.id,
231            "object": "chat.completion.chunk",
232            "created": self.created,
233            "model": self.model,
234            "choices": [choice],
235        });
236        format!("data: {value}\n\n")
237    }
238}
239
240// =============================================================================
241// Tests
242// =============================================================================
243
244#[cfg(test)]
245mod tests {
246    use super::*;
247    use crate::normalize::NormalizationErrorKind;
248
249    fn enc() -> SseEncoder {
250        SseEncoder::new("chatcmpl-1", "test-model", 1_729_000_000)
251    }
252
253    fn parse_data_frame(out: &str) -> serde_json::Value {
254        let line = out.lines().next().expect("non-empty output");
255        let payload = line.strip_prefix("data: ").expect("data: prefix");
256        serde_json::from_str(payload).expect("valid JSON")
257    }
258
259    #[test]
260    fn text_delta_encodes_to_content_chunk() {
261        let out = enc()
262            .encode(&LlmStreamEvent::TextDelta {
263                content: "hello".to_owned(),
264            })
265            .expect("frame");
266        assert!(out.starts_with("data: "));
267        assert!(out.ends_with("\n\n"));
268        let v = parse_data_frame(&out);
269        assert_eq!(v["object"], "chat.completion.chunk");
270        assert_eq!(v["id"], "chatcmpl-1");
271        assert_eq!(v["model"], "test-model");
272        assert_eq!(v["choices"][0]["delta"]["content"], "hello");
273        assert!(v["choices"][0]["finish_reason"].is_null());
274    }
275
276    #[test]
277    fn reasoning_delta_encodes_to_reasoning_content_chunk() {
278        let out = enc()
279            .encode(&LlmStreamEvent::ReasoningDelta {
280                content: "think".to_owned(),
281            })
282            .expect("frame");
283        let v = parse_data_frame(&out);
284        assert_eq!(v["choices"][0]["delta"]["reasoning_content"], "think");
285    }
286
287    #[test]
288    fn tool_call_delta_first_frame_includes_id_and_type() {
289        let out = enc()
290            .encode(&LlmStreamEvent::ToolCallDelta {
291                index: 0,
292                id: Some("tc1".to_owned()),
293                name: Some("search".to_owned()),
294                arguments: Some(r#"{"q":"r"}"#.to_owned()),
295            })
296            .expect("frame");
297        let v = parse_data_frame(&out);
298        let tc = &v["choices"][0]["delta"]["tool_calls"][0];
299        assert_eq!(tc["index"], 0);
300        assert_eq!(tc["id"], "tc1");
301        assert_eq!(tc["type"], "function");
302        assert_eq!(tc["function"]["name"], "search");
303        assert_eq!(tc["function"]["arguments"], r#"{"q":"r"}"#);
304    }
305
306    #[test]
307    fn tool_call_delta_continuation_omits_id_and_type() {
308        let out = enc()
309            .encode(&LlmStreamEvent::ToolCallDelta {
310                index: 0,
311                id: None,
312                name: None,
313                arguments: Some("more".to_owned()),
314            })
315            .expect("frame");
316        let v = parse_data_frame(&out);
317        let tc = &v["choices"][0]["delta"]["tool_calls"][0];
318        assert!(tc.get("id").is_none(), "id must be omitted on continuation");
319        assert!(
320            tc.get("type").is_none(),
321            "type must be omitted on continuation"
322        );
323        assert_eq!(tc["function"]["arguments"], "more");
324    }
325
326    #[test]
327    fn done_event_emits_only_finish_chunk_no_sentinel() {
328        let out = enc()
329            .encode(&LlmStreamEvent::Done {
330                finish_reason: "stop".to_owned(),
331            })
332            .expect("frame");
333        // Exactly one SSE frame -- [DONE] is the caller's responsibility now
334        // (see DONE_SENTINEL doc), since a trailing Usage event can
335        // legitimately follow Done.
336        let lines: Vec<&str> = out.lines().filter(|l| !l.is_empty()).collect();
337        assert_eq!(lines.len(), 1, "Done emits exactly one data: line");
338        let v: serde_json::Value =
339            serde_json::from_str(lines[0].strip_prefix("data: ").unwrap()).unwrap();
340        assert_eq!(v["choices"][0]["finish_reason"], "stop");
341    }
342
343    #[test]
344    fn usage_event_encodes_to_trailing_chunk_with_empty_choices() {
345        let out = enc()
346            .encode(&LlmStreamEvent::Usage {
347                prompt_tokens: 123,
348                completion_tokens: 45,
349                total_tokens: 168,
350                cached_tokens: None,
351            })
352            .expect("frame");
353        let v = parse_data_frame(&out);
354        assert_eq!(v["object"], "chat.completion.chunk");
355        assert_eq!(v["id"], "chatcmpl-1");
356        assert_eq!(v["model"], "test-model");
357        assert!(
358            v["choices"].as_array().is_some_and(Vec::is_empty),
359            "usage chunk must carry an empty choices array, not omit it"
360        );
361        assert_eq!(v["usage"]["prompt_tokens"], 123);
362        assert_eq!(v["usage"]["completion_tokens"], 45);
363        assert_eq!(v["usage"]["total_tokens"], 168);
364        assert!(
365            v["usage"].get("prompt_tokens_details").is_none(),
366            "an unreported cached-token count must not synthesize the details object"
367        );
368    }
369
370    /// A reported count is re-emitted under the OpenAI-standard nesting, so
371    /// clients (e.g. the Copilot LLM Gateway extension's `promptTokenDetails`)
372    /// see it exactly where they expect.
373    #[test]
374    fn usage_event_re_emits_a_reported_cached_token_count() {
375        let out = enc()
376            .encode(&LlmStreamEvent::Usage {
377                prompt_tokens: 123,
378                completion_tokens: 45,
379                total_tokens: 168,
380                cached_tokens: Some(100),
381            })
382            .expect("frame");
383        let v = parse_data_frame(&out);
384        assert_eq!(v["usage"]["prompt_tokens_details"]["cached_tokens"], 100);
385    }
386
387    /// Zero reused tokens is a real measurement, not a missing one, so it must
388    /// survive encoding rather than being elided like `None`.
389    #[test]
390    fn usage_event_distinguishes_zero_cached_tokens_from_absent() {
391        let out = enc()
392            .encode(&LlmStreamEvent::Usage {
393                prompt_tokens: 123,
394                completion_tokens: 45,
395                total_tokens: 168,
396                cached_tokens: Some(0),
397            })
398            .expect("frame");
399        let v = parse_data_frame(&out);
400        assert_eq!(v["usage"]["prompt_tokens_details"]["cached_tokens"], 0);
401    }
402
403    #[test]
404    fn upstream_error_event_encodes_to_bare_error_frame_no_sentinel() {
405        let out = enc()
406            .encode(&LlmStreamEvent::UpstreamError {
407                message: "Context window limit reached.".to_owned(),
408                error_type: "context_length_exceeded".to_owned(),
409                code: "context_length_exceeded".to_owned(),
410            })
411            .expect("frame");
412        let lines: Vec<&str> = out.lines().filter(|l| !l.is_empty()).collect();
413        assert_eq!(lines.len(), 1, "expects only the bare error frame");
414        let v: serde_json::Value =
415            serde_json::from_str(lines[0].strip_prefix("data: ").unwrap()).unwrap();
416        assert_eq!(v["error"]["message"], "Context window limit reached.");
417        assert_eq!(v["error"]["type"], "context_length_exceeded");
418        assert_eq!(v["error"]["code"], "context_length_exceeded");
419        assert!(
420            v.get("choices").is_none(),
421            "inline error frame must not carry a choices key at all"
422        );
423        assert!(
424            v.get("id").is_none(),
425            "inline error frame is deliberately bare, no envelope fields"
426        );
427    }
428
429    #[test]
430    fn prompt_progress_encodes_to_top_level_field() {
431        let out = enc()
432            .encode(&LlmStreamEvent::PromptProgress {
433                processed: 2,
434                total: 8,
435                cached: 1,
436                time_ms: 100,
437            })
438            .expect("frame");
439        let v = parse_data_frame(&out);
440        assert_eq!(v["prompt_progress"]["processed"], 2);
441        assert_eq!(v["prompt_progress"]["total"], 8);
442        assert_eq!(v["prompt_progress"]["cache"], 1);
443        assert_eq!(v["prompt_progress"]["time_ms"], 100);
444        assert!(v.get("choices").is_none());
445    }
446
447    #[test]
448    fn normalization_error_is_suppressed() {
449        let out = enc().encode(&LlmStreamEvent::NormalizationError {
450            kind: NormalizationErrorKind::MalformedToolCallJson {
451                raw: "<tool_call>oops".to_owned(),
452            },
453            raw: "<tool_call>oops".to_owned(),
454        });
455        assert!(
456            out.is_none(),
457            "NormalizationError must never reach the wire"
458        );
459    }
460}