1use serde_json::{Value, json};
29
30use crate::LlmStreamEvent;
31
32pub const DONE_SENTINEL: &str = "data: [DONE]\n\n";
41
42#[derive(Debug, Clone)]
47pub struct SseEncoder {
48 pub id: String,
50 pub model: String,
52 pub created: u64,
54}
55
56impl SseEncoder {
57 #[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 #[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 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 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 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 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 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 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#[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 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 #[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 #[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}