1use crate::tasks::generate::ToolCall;
7use crate::TokenUsage;
8use std::collections::HashMap;
9
10#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
18pub enum StreamEvent {
19 TextDelta(String),
21 ToolCallStart {
23 name: String,
24 index: usize,
25 id: Option<String>,
26 },
27 ToolCallDelta {
29 index: usize,
30 arguments_delta: String,
31 },
32 ToolCallArguments {
36 index: usize,
37 arguments: String,
38 authoritative: bool,
39 },
40 Usage {
48 input_tokens: u64,
53 output_tokens: u64,
54 cache_read_input_tokens: u64,
58 cache_creation_input_tokens: u64,
62 },
63 StopReason(String),
69 ProviderCost(f64),
74 ProviderOutputItem(serde_json::Value),
79 Error(String),
83 Attribution(crate::outcome::UnattributedReason),
89 Substituted(StreamSubstitution),
98 Done {
100 text: String,
101 tool_calls: Vec<ToolCall>,
102 },
103}
104
105#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
107pub struct StreamSubstitution {
108 pub model_id: String,
110 pub substituted_from: Option<String>,
113 pub substitution_reason: Option<crate::FallbackReason>,
115}
116
117impl StreamSubstitution {
118 pub fn annotate(&self, frame: &mut serde_json::Value) {
123 frame["local_last_resort"] = serde_json::json!(true);
124 if let Some(from) = &self.substituted_from {
125 frame["substituted_from"] = serde_json::json!(from);
126 }
127 if let Some(reason) = self.substitution_reason {
128 frame["substitution_reason"] = serde_json::json!(reason);
129 }
130 }
131}
132
133fn responses_usage_event(resp: &serde_json::Value) -> Option<StreamEvent> {
138 let u = resp.get("usage")?;
139 let input_total = u.get("input_tokens").and_then(|v| v.as_u64()).unwrap_or(0);
140 let cached = u
141 .get("input_tokens_details")
142 .and_then(|d| d.get("cached_tokens"))
143 .and_then(|v| v.as_u64())
144 .unwrap_or(0)
145 .min(input_total);
146 Some(StreamEvent::Usage {
147 input_tokens: input_total - cached,
148 output_tokens: u.get("output_tokens").and_then(|v| v.as_u64()).unwrap_or(0),
149 cache_read_input_tokens: cached,
150 cache_creation_input_tokens: 0,
151 })
152}
153
154pub fn parse_openai_responses_sse_line(event_type: &str, data: &str) -> Vec<StreamEvent> {
162 let json: serde_json::Value = match serde_json::from_str(data) {
163 Ok(v) => v,
164 Err(_) => return Vec::new(),
165 };
166 let mut events = Vec::new();
167 match event_type {
168 "response.output_text.delta" => {
169 if let Some(d) = json.get("delta").and_then(|d| d.as_str()) {
170 if !d.is_empty() {
171 events.push(StreamEvent::TextDelta(d.to_string()));
172 }
173 }
174 }
175 "response.output_item.added" | "response.output_item.done" => {
176 if let Some(item) = json.get("item") {
177 if item.get("type").and_then(|t| t.as_str()) == Some("function_call") {
178 let name = item
179 .get("name")
180 .and_then(|n| n.as_str())
181 .unwrap_or("")
182 .to_string();
183 let id = item
184 .get("call_id")
185 .or_else(|| item.get("id"))
186 .and_then(|i| i.as_str())
187 .map(|s| s.to_string());
188 let index = json
189 .get("output_index")
190 .and_then(|v| v.as_u64())
191 .unwrap_or(0) as usize;
192 if !name.is_empty() && event_type == "response.output_item.added" {
196 events.push(StreamEvent::ToolCallStart { name, index, id });
197 }
198 if let Some(arguments) = item
199 .get("arguments")
200 .and_then(|value| value.as_str())
201 .filter(|arguments| !arguments.is_empty())
202 {
203 events.push(StreamEvent::ToolCallArguments {
204 index,
205 arguments: arguments.to_string(),
206 authoritative: event_type == "response.output_item.done",
207 });
208 }
209 } else if event_type == "response.output_item.done"
210 && item.get("type").and_then(|value| value.as_str()) == Some("reasoning")
211 {
212 events.push(StreamEvent::ProviderOutputItem(item.clone()));
213 }
214 }
215 }
216 "response.function_call_arguments.delta" => {
217 if let Some(d) = json.get("delta").and_then(|d| d.as_str()) {
218 let index = json
219 .get("output_index")
220 .and_then(|v| v.as_u64())
221 .unwrap_or(0) as usize;
222 events.push(StreamEvent::ToolCallDelta {
223 index,
224 arguments_delta: d.to_string(),
225 });
226 }
227 }
228 "response.function_call_arguments.done" => {
229 if let Some(arguments) = json.get("arguments").and_then(|value| value.as_str()) {
230 let index = json
231 .get("output_index")
232 .and_then(|value| value.as_u64())
233 .unwrap_or(0) as usize;
234 events.push(StreamEvent::ToolCallArguments {
235 index,
236 arguments: arguments.to_string(),
237 authoritative: true,
238 });
239 }
240 }
241 "response.completed" | "response.incomplete" => {
242 if let Some(resp) = json.get("response") {
243 if let Some(usage) = responses_usage_event(resp) {
244 events.push(usage);
245 }
246 if event_type == "response.incomplete" {
247 let reason = resp
248 .pointer("/incomplete_details/reason")
249 .and_then(|r| r.as_str())
250 .unwrap_or("incomplete");
251 events.push(StreamEvent::StopReason(reason.to_string()));
252 events.push(StreamEvent::Error(format!(
253 "managed inference incomplete: {reason}"
254 )));
255 } else {
256 events.push(StreamEvent::Done {
261 text: String::new(),
262 tool_calls: Vec::new(),
263 });
264 }
265 }
266 }
267 "error" | "response.failed" => {
268 if let Some(usage) = json.get("response").and_then(responses_usage_event) {
271 events.push(usage);
272 }
273 let pick = |field: &str| {
274 json.pointer(&format!("/error/{field}"))
275 .or_else(|| json.pointer(&format!("/response/error/{field}")))
276 .and_then(|value| value.as_str())
277 .map(str::trim)
278 .filter(|value| !value.is_empty())
279 .map(str::to_string)
280 };
281 let message = pick("message");
282 let kind = pick("type");
296 let code = pick("code");
297 let mut detail = message.unwrap_or_else(|| "managed inference failed".to_string());
298 let tags: Vec<String> = [("type", kind), ("code", code)]
299 .into_iter()
300 .filter_map(|(label, value)| value.map(|v| format!("{label}={v}")))
301 .collect();
302 if !tags.is_empty() {
303 detail.push_str(&format!(" ({})", tags.join(", ")));
304 }
305 events.push(StreamEvent::Error(detail));
309 }
310 _ => {}
311 }
312 events
313}
314
315pub fn error_tags(detail: &str) -> (Option<&str>, Option<&str>) {
327 let Some(open) = detail.rfind(" (") else {
328 return (None, None);
329 };
330 let Some(close) = detail[open..].rfind(')') else {
331 return (None, None);
332 };
333 let mut kind = None;
334 let mut code = None;
335 for part in detail[open + 2..open + close].split(", ") {
336 if let Some(v) = part.strip_prefix("type=") {
337 kind = Some(v);
338 } else if let Some(v) = part.strip_prefix("code=") {
339 code = Some(v);
340 }
341 }
342 (kind, code)
343}
344
345pub fn content_refusal_tags(detail: &str) -> Option<(Option<String>, Option<String>)> {
369 let (kind, code) = error_tags(detail);
370 let refused = |v: &str| {
371 let v = v.to_ascii_lowercase();
372 v == "refusal"
374 || v.contains("content_policy")
375 || v.contains("content_filter")
376 || v.contains("moderation")
377 || v.contains("safety")
378 };
379 (code.is_some_and(refused) || kind.is_some_and(refused))
380 .then(|| (kind.map(str::to_string), code.map(str::to_string)))
381}
382
383pub(crate) fn google_output_tokens(usage: &serde_json::Value) -> u64 {
392 let count = |field: &str| usage.get(field).and_then(|v| v.as_u64()).unwrap_or(0);
393 count("candidatesTokenCount") + count("thoughtsTokenCount")
394}
395
396pub fn parse_google_sse_line(data: &str) -> Vec<StreamEvent> {
404 let json: serde_json::Value = match serde_json::from_str(data) {
405 Ok(v) => v,
406 Err(_) => return Vec::new(),
407 };
408 let mut events = Vec::new();
409 if let Some(parts) = json
410 .pointer("/candidates/0/content/parts")
411 .and_then(|p| p.as_array())
412 {
413 for (i, part) in parts.iter().enumerate() {
414 if let Some(t) = part.get("text").and_then(|t| t.as_str()) {
415 if !t.is_empty() {
416 events.push(StreamEvent::TextDelta(t.to_string()));
417 }
418 }
419 if let Some(fc) = part.get("functionCall") {
420 let name = fc
421 .get("name")
422 .and_then(|n| n.as_str())
423 .unwrap_or("")
424 .to_string();
425 if !name.is_empty() {
426 let args = fc
427 .get("args")
428 .map(|a| a.to_string())
429 .unwrap_or_else(|| "{}".to_string());
430 events.push(StreamEvent::ToolCallStart {
431 name,
432 index: i,
433 id: None,
434 });
435 events.push(StreamEvent::ToolCallDelta {
436 index: i,
437 arguments_delta: args,
438 });
439 }
440 }
441 }
442 }
443 if let Some(u) = json.get("usageMetadata") {
444 events.push(StreamEvent::Usage {
445 input_tokens: u
446 .get("promptTokenCount")
447 .and_then(|v| v.as_u64())
448 .unwrap_or(0),
449 output_tokens: google_output_tokens(u),
450 cache_read_input_tokens: 0,
452 cache_creation_input_tokens: 0,
453 });
454 }
455 if let Some(fr) = json
456 .pointer("/candidates/0/finishReason")
457 .and_then(|r| r.as_str())
458 {
459 events.push(StreamEvent::StopReason(fr.to_string()));
460 }
461 events
462}
463
464pub(crate) fn fold_openai_chat_sse(body: &str) -> Result<String, crate::InferenceError> {
477 if body.trim_start().starts_with('{') {
479 return Ok(body.to_string());
480 }
481 let mut chunks = 0usize;
482 let mut done = false;
483 let mut content = String::new();
484 let mut reasoning = String::new();
485 let mut tools: Vec<(u64, String, String, String)> = Vec::new();
487 let mut finish_reason = serde_json::Value::Null;
488 let mut usage = serde_json::Value::Null;
489 for line in body.lines() {
490 let Some(data) = line.strip_prefix("data:").map(str::trim) else {
491 continue;
492 };
493 if data == "[DONE]" {
494 done = true;
495 continue;
496 }
497 if data.is_empty() {
498 continue;
499 }
500 let chunk: serde_json::Value = serde_json::from_str(data).map_err(|e| {
501 crate::InferenceError::InferenceFailed(format!("parse stream chunk: {e}"))
502 })?;
503 chunks += 1;
504 if let Some(error) = chunk.get("error") {
505 let message = error
506 .get("message")
507 .and_then(|m| m.as_str())
508 .unwrap_or("the server reported an error mid-stream")
509 .to_string();
510 let status = error.get("code").and_then(|c| c.as_u64()).map(|c| c as u16);
514 return Err(match status {
515 Some(code) if code == 429 || code >= 500 => crate::InferenceError::Transient {
516 status: Some(code),
517 message,
518 },
519 _ => crate::InferenceError::InferenceFailed(message),
520 });
521 }
522 if let Some(u) = chunk.get("usage").filter(|u| !u.is_null()) {
523 usage = u.clone();
524 }
525 let Some(choice) = chunk.pointer("/choices/0") else {
526 continue;
527 };
528 if let Some(fr) = choice.get("finish_reason").filter(|f| !f.is_null()) {
529 finish_reason = fr.clone();
530 }
531 let Some(delta) = choice.get("delta") else {
532 continue;
533 };
534 if let Some(t) = delta.get("content").and_then(|c| c.as_str()) {
535 content.push_str(t);
536 }
537 if let Some(t) = delta
538 .get("reasoning_content")
539 .or_else(|| delta.get("reasoning"))
540 .and_then(|c| c.as_str())
541 {
542 reasoning.push_str(t);
543 }
544 for call in delta
545 .get("tool_calls")
546 .and_then(|t| t.as_array())
547 .into_iter()
548 .flatten()
549 {
550 let index = call.get("index").and_then(|i| i.as_u64()).unwrap_or(0);
551 let slot = match tools.iter().position(|t| t.0 == index) {
552 Some(i) => &mut tools[i],
553 None => {
554 tools.push((index, String::new(), String::new(), String::new()));
555 tools.last_mut().expect("just pushed")
556 }
557 };
558 if let Some(id) = call.get("id").and_then(|v| v.as_str()) {
559 slot.1 = id.to_string();
560 }
561 if let Some(name) = call.pointer("/function/name").and_then(|v| v.as_str()) {
562 slot.2.push_str(name);
563 }
564 if let Some(args) = call.pointer("/function/arguments").and_then(|v| v.as_str()) {
565 slot.3.push_str(args);
566 }
567 }
568 }
569 if chunks == 0 || (!done && finish_reason.is_null()) {
572 return Err(crate::InferenceError::InferenceFailed(
573 "the stream ended before the server finished the response".to_string(),
574 ));
575 }
576 let mut message = serde_json::json!({ "role": "assistant", "content": content });
577 if !reasoning.is_empty() {
578 message["reasoning_content"] = serde_json::Value::String(reasoning);
579 }
580 if !tools.is_empty() {
581 message["tool_calls"] = tools
582 .into_iter()
583 .map(|(_, id, name, arguments)| {
584 serde_json::json!({
585 "id": id,
586 "type": "function",
587 "function": { "name": name, "arguments": arguments },
588 })
589 })
590 .collect();
591 }
592 let mut response = serde_json::json!({
593 "choices": [{ "index": 0, "message": message, "finish_reason": finish_reason }],
594 });
595 if !usage.is_null() {
597 response["usage"] = usage;
598 }
599 Ok(response.to_string())
600}
601
602pub fn parse_openai_sse_line(line: &str) -> Vec<StreamEvent> {
605 let data = match line.strip_prefix("data: ") {
606 Some(d) => d,
607 None => return Vec::new(),
608 };
609 if data == "[DONE]" {
610 return Vec::new();
611 }
612
613 let json: serde_json::Value = match serde_json::from_str(data) {
614 Ok(v) => v,
615 Err(_) => return Vec::new(),
616 };
617
618 let mut events = Vec::new();
619
620 if let Some(reason) = json
624 .get("choices")
625 .and_then(|c| c.as_array())
626 .and_then(|c| c.first())
627 .and_then(|c| c.get("finish_reason"))
628 .and_then(|r| r.as_str())
629 {
630 if !reason.is_empty() {
631 events.push(StreamEvent::StopReason(reason.to_string()));
632 }
633 }
634
635 if let Some(delta) = json
638 .get("choices")
639 .and_then(|c| c.as_array())
640 .and_then(|c| c.first())
641 .and_then(|c| c.get("delta"))
642 {
643 if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
644 if !content.is_empty() {
645 events.push(StreamEvent::TextDelta(content.to_string()));
646 }
647 }
648
649 if let Some(tool_calls) = delta.get("tool_calls").and_then(|t| t.as_array()) {
651 for tc in tool_calls {
652 let index = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
653 if let Some(function) = tc.get("function") {
654 if let Some(name) = function.get("name").and_then(|n| n.as_str()) {
655 let id = tc.get("id").and_then(|i| i.as_str()).map(|s| s.to_string());
656 events.push(StreamEvent::ToolCallStart {
657 name: name.to_string(),
658 index,
659 id,
660 });
661 }
662 if let Some(args) = function.get("arguments").and_then(|a| a.as_str()) {
663 if !args.is_empty() {
664 events.push(StreamEvent::ToolCallDelta {
665 index,
666 arguments_delta: args.to_string(),
667 });
668 }
669 }
670 }
671 }
672 }
673 }
674
675 if let Some(usage) = json.get("usage") {
679 let input = usage
680 .get("prompt_tokens")
681 .and_then(|n| n.as_u64())
682 .unwrap_or(0);
683 let output = usage
684 .get("completion_tokens")
685 .and_then(|n| n.as_u64())
686 .unwrap_or(0);
687 let cached = usage
690 .get("prompt_tokens_details")
691 .and_then(|d| d.get("cached_tokens"))
692 .and_then(|n| n.as_u64())
693 .unwrap_or(0)
694 .min(input);
695 if input != 0 || output != 0 {
696 events.push(StreamEvent::Usage {
697 input_tokens: input - cached,
698 output_tokens: output,
699 cache_read_input_tokens: cached,
700 cache_creation_input_tokens: 0,
701 });
702 }
703 }
704
705 events
706}
707
708pub fn parse_anthropic_sse_line(event_type: &str, data: &str) -> Vec<StreamEvent> {
710 match event_type {
711 "content_block_delta" => {
712 let json: serde_json::Value = match serde_json::from_str(data) {
713 Ok(v) => v,
714 Err(_) => return Vec::new(),
715 };
716 let delta = match json.get("delta") {
717 Some(d) => d,
718 None => return Vec::new(),
719 };
720 let delta_type = match delta.get("type").and_then(|t| t.as_str()) {
721 Some(t) => t,
722 None => return Vec::new(),
723 };
724
725 match delta_type {
726 "text_delta" => match delta.get("text").and_then(|t| t.as_str()) {
727 Some(text) => vec![StreamEvent::TextDelta(text.to_string())],
728 None => Vec::new(),
729 },
730 "input_json_delta" => match delta.get("partial_json").and_then(|p| p.as_str()) {
731 Some(partial) => {
732 let index =
733 json.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
734 vec![StreamEvent::ToolCallDelta {
735 index,
736 arguments_delta: partial.to_string(),
737 }]
738 }
739 None => Vec::new(),
740 },
741 _ => Vec::new(),
742 }
743 }
744 "content_block_start" => {
745 let json: serde_json::Value = match serde_json::from_str(data) {
746 Ok(v) => v,
747 Err(_) => return Vec::new(),
748 };
749 let block = match json.get("content_block") {
750 Some(b) => b,
751 None => return Vec::new(),
752 };
753 if block.get("type").and_then(|t| t.as_str()) == Some("tool_use") {
754 if let Some(name) = block.get("name").and_then(|n| n.as_str()) {
755 let index = json.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
756 let id = block
757 .get("id")
758 .and_then(|i| i.as_str())
759 .map(|s| s.to_string());
760 return vec![StreamEvent::ToolCallStart {
761 name: name.to_string(),
762 index,
763 id,
764 }];
765 }
766 }
767 Vec::new()
768 }
769 "message_start" => {
773 let json: serde_json::Value = match serde_json::from_str(data) {
774 Ok(v) => v,
775 Err(_) => return Vec::new(),
776 };
777 let Some(usage) = json.pointer("/message/usage") else {
778 return Vec::new();
779 };
780 let input = usage
781 .get("input_tokens")
782 .and_then(|n| n.as_u64())
783 .unwrap_or(0);
784 let output = usage
785 .get("output_tokens")
786 .and_then(|n| n.as_u64())
787 .unwrap_or(0);
788 let cache_read = usage
791 .get("cache_read_input_tokens")
792 .and_then(|n| n.as_u64())
793 .unwrap_or(0);
794 let cache_creation = usage
795 .get("cache_creation_input_tokens")
796 .and_then(|n| n.as_u64())
797 .unwrap_or(0);
798 if input == 0 && output == 0 && cache_read == 0 && cache_creation == 0 {
799 return Vec::new();
800 }
801 vec![StreamEvent::Usage {
802 input_tokens: input,
803 output_tokens: output,
804 cache_read_input_tokens: cache_read,
805 cache_creation_input_tokens: cache_creation,
806 }]
807 }
808 "message_delta" => {
812 let json: serde_json::Value = match serde_json::from_str(data) {
813 Ok(v) => v,
814 Err(_) => return Vec::new(),
815 };
816 let mut events = Vec::new();
817 if let Some(reason) = json.pointer("/delta/stop_reason").and_then(|r| r.as_str()) {
820 if !reason.is_empty() {
821 events.push(StreamEvent::StopReason(reason.to_string()));
822 }
823 }
824 if let Some(usage) = json.get("usage") {
825 let input = usage
826 .get("input_tokens")
827 .and_then(|n| n.as_u64())
828 .unwrap_or(0);
829 let output = usage
830 .get("output_tokens")
831 .and_then(|n| n.as_u64())
832 .unwrap_or(0);
833 if input != 0 || output != 0 {
834 events.push(StreamEvent::Usage {
835 input_tokens: input,
836 output_tokens: output,
837 cache_read_input_tokens: 0,
840 cache_creation_input_tokens: 0,
841 });
842 }
843 }
844 events
845 }
846 _ => Vec::new(),
847 }
848}
849
850pub fn enforce_stop_sequences(
871 mut source: tokio::sync::mpsc::Receiver<StreamEvent>,
872 stops: Vec<String>,
873) -> tokio::sync::mpsc::Receiver<StreamEvent> {
874 let stops: Vec<String> = stops.into_iter().filter(|stop| !stop.is_empty()).collect();
875 if stops.is_empty() {
876 return source;
877 }
878 let (tx, rx) = tokio::sync::mpsc::channel(64);
879 tokio::spawn(async move {
880 let mut filter = StopSequenceFilter::new(stops);
881 while let Some(event) = source.recv().await {
882 for out in filter.push(event) {
883 if tx.send(out).await.is_err() {
884 return;
885 }
886 }
887 }
888 for out in filter.finish() {
889 if tx.send(out).await.is_err() {
890 return;
891 }
892 }
893 });
894 rx
895}
896
897struct StopSequenceFilter {
900 stops: Vec<String>,
901 held: String,
904 stopped: bool,
906 reported_stop: bool,
908 provider_reason: Option<String>,
911}
912
913impl StopSequenceFilter {
914 fn new(stops: Vec<String>) -> Self {
915 Self {
916 stops,
917 held: String::new(),
918 stopped: false,
919 reported_stop: false,
920 provider_reason: None,
921 }
922 }
923
924 fn ambiguous_tail(&self) -> usize {
927 self.held
928 .char_indices()
929 .map(|(index, _)| index)
930 .find(|&index| {
931 let tail = &self.held[index..];
932 self.stops
933 .iter()
934 .any(|stop| stop.len() > tail.len() && stop.starts_with(tail))
935 })
936 .map_or(0, |index| self.held.len() - index)
937 }
938
939 fn take_held(&mut self) -> Vec<StreamEvent> {
940 if self.held.is_empty() {
941 Vec::new()
942 } else {
943 vec![StreamEvent::TextDelta(std::mem::take(&mut self.held))]
944 }
945 }
946
947 fn stop_reason(&mut self) -> StreamEvent {
948 self.reported_stop = true;
949 StreamEvent::StopReason(crate::tasks::generate::STOP_SEQUENCE_STOP_REASON.to_string())
950 }
951
952 fn push(&mut self, event: StreamEvent) -> Vec<StreamEvent> {
953 match event {
954 StreamEvent::TextDelta(delta) => {
955 if self.stopped {
956 return Vec::new();
957 }
958 self.held.push_str(&delta);
959 let earliest = self
960 .stops
961 .iter()
962 .filter_map(|stop| self.held.find(stop.as_str()))
963 .min();
964 if let Some(cut) = earliest {
965 self.held.truncate(cut);
966 self.stopped = true;
967 return self.take_held();
968 }
969 let emit = self.held.len() - self.ambiguous_tail();
970 if emit == 0 {
971 return Vec::new();
972 }
973 let rest = self.held.split_off(emit);
974 vec![StreamEvent::TextDelta(std::mem::replace(
975 &mut self.held,
976 rest,
977 ))]
978 }
979 StreamEvent::StopReason(reason) => {
980 self.provider_reason = Some(reason);
983 Vec::new()
984 }
985 StreamEvent::Done { text, tool_calls } => {
986 let mut out = self.take_held();
987 let cut = crate::tasks::generate::truncate_at_stop(&text, &self.stops);
988 if cut.len() < text.len() {
989 self.stopped = true;
990 }
991 out.extend(self.reason_event());
992 if !tool_calls.is_empty() {
993 tracing::warn!(
994 tool_calls = tool_calls.len(),
995 "stop sequences apply to streamed text only; tool-call arguments \
996 in this stream were not checked against them"
997 );
998 }
999 out.push(StreamEvent::Done {
1000 text: cut,
1001 tool_calls,
1002 });
1003 out
1004 }
1005 StreamEvent::Error(message) => {
1006 let mut out = self.take_held();
1008 out.extend(self.reason_event());
1009 out.push(StreamEvent::Error(message));
1010 out
1011 }
1012 other @ (StreamEvent::ToolCallStart { .. }
1016 | StreamEvent::ToolCallDelta { .. }
1017 | StreamEvent::ToolCallArguments { .. }
1018 | StreamEvent::Usage { .. }
1019 | StreamEvent::ProviderCost(_)
1020 | StreamEvent::ProviderOutputItem(_)
1021 | StreamEvent::Attribution(_)
1022 | StreamEvent::Substituted(_)) => vec![other],
1023 }
1024 }
1025
1026 fn reason_event(&mut self) -> Option<StreamEvent> {
1029 if self.reported_stop {
1030 return None;
1031 }
1032 if self.stopped {
1033 self.provider_reason = None;
1034 return Some(self.stop_reason());
1035 }
1036 self.provider_reason.take().map(|reason| {
1037 self.reported_stop = true;
1038 StreamEvent::StopReason(reason)
1039 })
1040 }
1041
1042 fn finish(&mut self) -> Vec<StreamEvent> {
1044 let mut out = self.take_held();
1045 out.extend(self.reason_event());
1046 out
1047 }
1048}
1049
1050#[derive(Default)]
1052pub struct StreamAccumulator {
1053 pub text: String,
1054 tool_names: HashMap<usize, String>,
1055 tool_args: HashMap<usize, String>,
1056 tool_args_fallback: HashMap<usize, String>,
1057 tool_args_authoritative: HashMap<usize, String>,
1058 tool_ids: HashMap<usize, String>,
1059 input_tokens: u64,
1064 output_tokens: u64,
1068 cache_read_input_tokens: u64,
1070 cache_creation_input_tokens: u64,
1072 saw_usage: bool,
1077 stop_reason: Option<String>,
1080 provider_output_items: Vec<serde_json::Value>,
1082 reported_cost_usd: Option<f64>,
1084}
1085
1086impl StreamAccumulator {
1087 pub fn push(&mut self, event: &StreamEvent) {
1088 match event {
1089 StreamEvent::TextDelta(t) => self.text.push_str(t),
1090 StreamEvent::ToolCallStart { name, index, id } => {
1091 self.tool_names.insert(*index, name.clone());
1092 self.tool_args.entry(*index).or_default();
1093 if let Some(id) = id {
1094 self.tool_ids.insert(*index, id.clone());
1095 }
1096 }
1097 StreamEvent::ToolCallDelta {
1098 index,
1099 arguments_delta,
1100 } => {
1101 self.tool_args
1102 .entry(*index)
1103 .or_default()
1104 .push_str(arguments_delta);
1105 }
1106 StreamEvent::ToolCallArguments {
1107 index,
1108 arguments,
1109 authoritative,
1110 } => {
1111 let destination = if *authoritative {
1112 &mut self.tool_args_authoritative
1113 } else {
1114 &mut self.tool_args_fallback
1115 };
1116 destination.insert(*index, arguments.clone());
1117 }
1118 StreamEvent::Usage {
1119 input_tokens,
1120 output_tokens,
1121 cache_read_input_tokens,
1122 cache_creation_input_tokens,
1123 } => {
1124 self.saw_usage = true;
1125 if *input_tokens > self.input_tokens {
1132 self.input_tokens = *input_tokens;
1133 }
1134 if *output_tokens > self.output_tokens {
1135 self.output_tokens = *output_tokens;
1136 }
1137 if *cache_read_input_tokens > self.cache_read_input_tokens {
1138 self.cache_read_input_tokens = *cache_read_input_tokens;
1139 }
1140 if *cache_creation_input_tokens > self.cache_creation_input_tokens {
1141 self.cache_creation_input_tokens = *cache_creation_input_tokens;
1142 }
1143 }
1144 StreamEvent::StopReason(reason) => {
1145 self.stop_reason = Some(reason.clone());
1146 }
1147 StreamEvent::ProviderOutputItem(item) => {
1148 self.provider_output_items.push(item.clone());
1149 }
1150 StreamEvent::ProviderCost(cost) => {
1151 self.reported_cost_usd = Some(*cost);
1152 }
1153 StreamEvent::Error(_) | StreamEvent::Attribution(_) | StreamEvent::Substituted(_) => {}
1154 StreamEvent::Done { .. } => {}
1155 }
1156 }
1157
1158 pub fn finish(self) -> (String, Vec<ToolCall>) {
1159 let (text, tool_calls, _, _) = self.finish_with_usage();
1160 (text, tool_calls)
1161 }
1162
1163 pub fn finish_with_usage(self) -> (String, Vec<ToolCall>, Option<TokenUsage>, Option<String>) {
1171 let (text, tool_calls, usage, stop_reason, _) = self.finish_with_provider_output_items();
1172 (text, tool_calls, usage, stop_reason)
1173 }
1174
1175 pub fn finish_with_provider_output_items(
1178 self,
1179 ) -> (
1180 String,
1181 Vec<ToolCall>,
1182 Option<TokenUsage>,
1183 Option<String>,
1184 Vec<serde_json::Value>,
1185 ) {
1186 let mut tool_calls = Vec::new();
1187 let mut indices: Vec<usize> = self.tool_names.keys().copied().collect();
1188 indices.sort();
1189
1190 for idx in indices {
1191 let id = self.tool_ids.get(&idx).cloned();
1192 let name = self.tool_names.get(&idx).cloned().unwrap_or_default();
1193 let args_str = self
1194 .tool_args_authoritative
1195 .get(&idx)
1196 .or_else(|| {
1197 self.tool_args
1198 .get(&idx)
1199 .filter(|arguments| !arguments.is_empty())
1200 })
1201 .or_else(|| self.tool_args_fallback.get(&idx))
1202 .cloned()
1203 .unwrap_or_default();
1204 let arguments: HashMap<String, serde_json::Value> = if args_str.trim().is_empty() {
1205 HashMap::new()
1206 } else {
1207 match serde_json::from_str(&args_str) {
1208 Ok(arguments) => arguments,
1209 Err(error) => {
1210 tracing::warn!(
1211 tool = %name,
1212 output_index = idx,
1213 %error,
1214 "tool-call arguments were not valid JSON; returning an empty argument map"
1215 );
1216 HashMap::new()
1217 }
1218 }
1219 };
1220 tool_calls.push(ToolCall {
1221 id,
1222 name,
1223 arguments,
1224 });
1225 }
1226
1227 let usage = if self.saw_usage {
1228 Some(TokenUsage {
1229 prompt_tokens: self.input_tokens,
1230 completion_tokens: self.output_tokens,
1231 total_tokens: self.input_tokens + self.output_tokens,
1232 context_window: 0,
1236 cache_read_input_tokens: self.cache_read_input_tokens,
1239 cache_creation_input_tokens: self.cache_creation_input_tokens,
1240 reported_cost_usd: self.reported_cost_usd,
1241 reasoning_tokens: None,
1243 })
1244 } else {
1245 None
1246 };
1247
1248 let (text, tag_calls) = crate::tasks::generate::parse_tool_calls(&self.text);
1254 let (text, tool_calls) = if tool_calls.is_empty() && !tag_calls.is_empty() {
1255 (text, tag_calls)
1256 } else {
1257 (text, tool_calls)
1259 };
1260
1261 (
1262 text,
1263 tool_calls,
1264 usage,
1265 self.stop_reason,
1266 self.provider_output_items,
1267 )
1268 }
1269}
1270
1271pub fn parse_sse_lines(chunk: &str) -> Vec<(String, String)> {
1274 let mut events = Vec::new();
1275 let mut current_event = String::new();
1276 let mut current_data = String::new();
1277
1278 for line in chunk.lines() {
1279 if let Some(rest) = line.strip_prefix("event: ") {
1280 current_event = rest.to_string();
1281 } else if let Some(rest) = line.strip_prefix("data: ") {
1282 current_data = rest.to_string();
1283 } else if line.is_empty() && !current_data.is_empty() {
1284 events.push((
1285 if current_event.is_empty() {
1286 "message".to_string()
1287 } else {
1288 current_event.clone()
1289 },
1290 current_data.clone(),
1291 ));
1292 current_event.clear();
1293 current_data.clear();
1294 }
1295 }
1296
1297 if !current_data.is_empty() {
1299 events.push((
1300 if current_event.is_empty() {
1301 "message".to_string()
1302 } else {
1303 current_event
1304 },
1305 current_data,
1306 ));
1307 }
1308
1309 events
1310}
1311
1312#[cfg(test)]
1313mod tests {
1314
1315 #[test]
1316 fn a_folded_stream_must_be_a_complete_stream() {
1317 let json = r#"{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}]}"#;
1319 assert_eq!(fold_openai_chat_sse(json).unwrap(), json);
1320 let cut = "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"par\"}}]}\n\n";
1322 assert!(fold_openai_chat_sse(cut).is_err());
1323 assert!(fold_openai_chat_sse("").is_err());
1324 let no_usage = "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n";
1326 let folded: serde_json::Value =
1327 serde_json::from_str(&fold_openai_chat_sse(no_usage).unwrap()).unwrap();
1328 assert!(folded.get("usage").is_none(), "{folded}");
1329 assert_eq!(
1330 folded
1331 .pointer("/choices/0/message/content")
1332 .and_then(|c| c.as_str()),
1333 Some("ok")
1334 );
1335 }
1336
1337 #[test]
1340 fn a_mid_stream_server_error_keeps_its_status() {
1341 let busy = "data: {\"error\":{\"message\":\"Internal server error\",\"code\":500}}\n\n";
1342 assert!(matches!(
1343 fold_openai_chat_sse(busy),
1344 Err(crate::InferenceError::Transient {
1345 status: Some(500),
1346 ..
1347 })
1348 ));
1349 let bad = "data: {\"error\":{\"message\":\"bad request\",\"code\":400}}\n\n";
1350 assert!(matches!(
1351 fold_openai_chat_sse(bad),
1352 Err(crate::InferenceError::InferenceFailed(_))
1353 ));
1354 }
1355
1356 use super::*;
1357
1358 #[test]
1366 fn accumulated_usage_and_stop_reason_survive_finish() {
1367 let mut acc = StreamAccumulator::default();
1368 acc.push(&StreamEvent::TextDelta("hello".into()));
1369 acc.push(&StreamEvent::Usage {
1370 input_tokens: 28,
1371 output_tokens: 5,
1372 cache_read_input_tokens: 0,
1373 cache_creation_input_tokens: 0,
1374 });
1375 acc.push(&StreamEvent::StopReason("length".into()));
1376
1377 let (text, _tools, usage, stop) = acc.finish_with_usage();
1378 assert_eq!(text, "hello");
1379 let usage = usage.expect("a reported Usage event must not be dropped");
1380 assert_eq!(usage.prompt_tokens, 28);
1381 assert_eq!(usage.completion_tokens, 5);
1382 assert_eq!(usage.total_tokens, 33);
1383 assert_eq!(
1384 stop.as_deref(),
1385 Some("length"),
1386 "the provider stop_reason feeds was_truncated and was being dropped too"
1387 );
1388 }
1389
1390 #[test]
1393 fn absent_usage_stays_none_rather_than_zero() {
1394 let mut acc = StreamAccumulator::default();
1395 acc.push(&StreamEvent::TextDelta("hi".into()));
1396 let (_text, _tools, usage, stop) = acc.finish_with_usage();
1397 assert!(
1398 usage.is_none(),
1399 "no Usage event must yield None so callers can fall back to an estimator"
1400 );
1401 assert!(stop.is_none());
1402 }
1403
1404 #[test]
1412 fn managed_error_events_carry_type_and_code() {
1413 let events = parse_openai_responses_sse_line(
1414 "error",
1415 r#"{"error":{"message":"content refused","type":"invalid_request_error","code":"content_policy_violation"}}"#,
1416 );
1417 let StreamEvent::Error(msg) = events.first().expect("an error event") else {
1418 panic!("expected StreamEvent::Error, got {:?}", events.first());
1419 };
1420 assert!(msg.contains("content refused"), "message dropped: {msg}");
1421 assert!(
1422 msg.contains("type=invalid_request_error"),
1423 "type dropped: {msg}"
1424 );
1425 assert!(
1426 msg.contains("code=content_policy_violation"),
1427 "code dropped: {msg}"
1428 );
1429 }
1430
1431 #[test]
1435 fn a_messageless_managed_error_still_reports_its_code() {
1436 let events = parse_openai_responses_sse_line(
1437 "response.failed",
1438 r#"{"response":{"error":{"code":"content_filter"}}}"#,
1439 );
1440 let StreamEvent::Error(msg) = events.first().expect("an error event") else {
1441 panic!("expected StreamEvent::Error");
1442 };
1443 assert!(msg.contains("managed inference failed"), "{msg}");
1444 assert!(
1445 msg.contains("code=content_filter"),
1446 "classification lost: {msg}"
1447 );
1448 }
1449
1450 #[test]
1453 fn a_bare_managed_error_is_unchanged() {
1454 let events = parse_openai_responses_sse_line("error", r#"{"error":{}}"#);
1455 let StreamEvent::Error(msg) = events.first().expect("an error event") else {
1456 panic!("expected StreamEvent::Error");
1457 };
1458 assert_eq!(msg, "managed inference failed");
1459 }
1460
1461 #[test]
1462 fn parse_openai_text_delta() {
1463 let line = r#"data: {"choices":[{"delta":{"content":"Hello"}}]}"#;
1464 let events = parse_openai_sse_line(line);
1465 assert_eq!(events.len(), 1);
1466 match &events[0] {
1467 StreamEvent::TextDelta(t) => assert_eq!(t, "Hello"),
1468 other => panic!("expected TextDelta, got {:?}", other),
1469 }
1470 }
1471
1472 #[test]
1473 fn parse_openai_tool_call_start() {
1474 let line = r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"edit_file"}}]}}]}"#;
1475 let events = parse_openai_sse_line(line);
1476 assert_eq!(events.len(), 1);
1477 match &events[0] {
1478 StreamEvent::ToolCallStart { name, index, .. } => {
1479 assert_eq!(name, "edit_file");
1480 assert_eq!(*index, 0);
1481 }
1482 other => panic!("expected ToolCallStart, got {:?}", other),
1483 }
1484 }
1485
1486 #[test]
1487 fn parse_openai_tool_call_delta() {
1488 let line = r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"path\":"}}]}}]}"#;
1489 let events = parse_openai_sse_line(line);
1490 assert_eq!(events.len(), 1);
1491 match &events[0] {
1492 StreamEvent::ToolCallDelta {
1493 index,
1494 arguments_delta,
1495 } => {
1496 assert_eq!(*index, 0);
1497 assert!(arguments_delta.contains("path"));
1498 }
1499 other => panic!("expected ToolCallDelta, got {:?}", other),
1500 }
1501 }
1502
1503 #[test]
1504 fn parse_openai_multiple_tool_calls_in_chunk() {
1505 let line = r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"read_file"}},{"index":1,"function":{"name":"search"}}]}}]}"#;
1507 let events = parse_openai_sse_line(line);
1508 assert_eq!(events.len(), 2);
1509 match &events[0] {
1510 StreamEvent::ToolCallStart { name, index, .. } => {
1511 assert_eq!(name, "read_file");
1512 assert_eq!(*index, 0);
1513 }
1514 other => panic!("expected ToolCallStart, got {:?}", other),
1515 }
1516 match &events[1] {
1517 StreamEvent::ToolCallStart { name, index, .. } => {
1518 assert_eq!(name, "search");
1519 assert_eq!(*index, 1);
1520 }
1521 other => panic!("expected ToolCallStart, got {:?}", other),
1522 }
1523 }
1524
1525 #[test]
1526 fn parse_openai_done() {
1527 assert!(parse_openai_sse_line("data: [DONE]").is_empty());
1528 }
1529
1530 #[test]
1531 fn parse_anthropic_text_delta() {
1532 let data = r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"world"}}"#;
1533 let events = parse_anthropic_sse_line("content_block_delta", data);
1534 assert_eq!(events.len(), 1);
1535 match &events[0] {
1536 StreamEvent::TextDelta(t) => assert_eq!(t, "world"),
1537 other => panic!("expected TextDelta, got {:?}", other),
1538 }
1539 }
1540
1541 #[test]
1542 fn parse_anthropic_tool_start() {
1543 let data = r#"{"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"t1","name":"search","input":{}}}"#;
1544 let events = parse_anthropic_sse_line("content_block_start", data);
1545 assert_eq!(events.len(), 1);
1546 match &events[0] {
1547 StreamEvent::ToolCallStart { name, index, .. } => {
1548 assert_eq!(name, "search");
1549 assert_eq!(*index, 1);
1550 }
1551 other => panic!("expected ToolCallStart, got {:?}", other),
1552 }
1553 }
1554
1555 #[test]
1556 fn accumulator_builds_result() {
1557 let mut acc = StreamAccumulator::default();
1558 acc.push(&StreamEvent::TextDelta("Hello ".into()));
1559 acc.push(&StreamEvent::TextDelta("world".into()));
1560 acc.push(&StreamEvent::ToolCallStart {
1561 name: "search".into(),
1562 index: 0,
1563 id: None,
1564 });
1565 acc.push(&StreamEvent::ToolCallDelta {
1566 index: 0,
1567 arguments_delta: r#"{"q":"test"}"#.into(),
1568 });
1569
1570 let (text, tools) = acc.finish();
1571 assert_eq!(text, "Hello world");
1572 assert_eq!(tools.len(), 1);
1573 assert_eq!(tools[0].name, "search");
1574 assert!(tools[0].arguments.contains_key("q"));
1575 }
1576
1577 #[test]
1578 fn parse_sse_lines_openai_format() {
1579 let chunk = "data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\ndata: [DONE]\n\n";
1580 let events = parse_sse_lines(chunk);
1581 assert_eq!(events.len(), 2);
1582 assert_eq!(events[0].0, "message");
1583 assert_eq!(events[1].1, "[DONE]");
1584 }
1585
1586 #[test]
1587 fn parse_sse_lines_anthropic_format() {
1588 let chunk = "event: content_block_delta\ndata: {\"delta\":{\"type\":\"text_delta\",\"text\":\"Hi\"}}\n\n";
1589 let events = parse_sse_lines(chunk);
1590 assert_eq!(events.len(), 1);
1591 assert_eq!(events[0].0, "content_block_delta");
1592 }
1593
1594 #[test]
1595 fn parse_anthropic_message_start_emits_usage() {
1596 let data = r#"{"type":"message_start","message":{"id":"msg_1","role":"assistant","usage":{"input_tokens":245,"output_tokens":1}}}"#;
1597 let events = parse_anthropic_sse_line("message_start", data);
1598 assert_eq!(events.len(), 1);
1599 match &events[0] {
1600 StreamEvent::Usage {
1601 input_tokens,
1602 output_tokens,
1603 ..
1604 } => {
1605 assert_eq!(*input_tokens, 245);
1606 assert_eq!(*output_tokens, 1);
1607 }
1608 other => panic!("expected Usage, got {:?}", other),
1609 }
1610 }
1611
1612 #[test]
1613 fn parse_anthropic_message_delta_emits_stop_reason_and_usage() {
1614 let data = r#"{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":87}}"#;
1615 let events = parse_anthropic_sse_line("message_delta", data);
1616 assert_eq!(events.len(), 2);
1619 match &events[0] {
1620 StreamEvent::StopReason(reason) => assert_eq!(reason, "end_turn"),
1621 other => panic!("expected StopReason, got {:?}", other),
1622 }
1623 match &events[1] {
1624 StreamEvent::Usage {
1625 input_tokens,
1626 output_tokens,
1627 ..
1628 } => {
1629 assert_eq!(*input_tokens, 0);
1630 assert_eq!(*output_tokens, 87);
1631 }
1632 other => panic!("expected Usage, got {:?}", other),
1633 }
1634 }
1635
1636 #[test]
1637 fn parse_anthropic_message_delta_max_tokens_stop_reason() {
1638 let data = r#"{"type":"message_delta","delta":{"stop_reason":"max_tokens"},"usage":{"output_tokens":4096}}"#;
1641 let events = parse_anthropic_sse_line("message_delta", data);
1642 assert!(matches!(
1643 &events[0],
1644 StreamEvent::StopReason(r) if r == "max_tokens"
1645 ));
1646 }
1647
1648 #[test]
1649 fn parse_openai_finish_reason_length_surfaces() {
1650 let line = r#"data: {"choices":[{"delta":{"content":""},"finish_reason":"length"}]}"#;
1653 let events = parse_openai_sse_line(line);
1654 assert!(events
1655 .iter()
1656 .any(|e| matches!(e, StreamEvent::StopReason(r) if r == "length")));
1657 }
1658
1659 #[test]
1660 fn accumulator_captures_stop_reason() {
1661 let mut acc = StreamAccumulator::default();
1662 acc.push(&StreamEvent::TextDelta("partial".into()));
1663 acc.push(&StreamEvent::StopReason("max_tokens".into()));
1664 let (_, _, _, stop) = acc.finish_with_usage();
1665 assert_eq!(stop.as_deref(), Some("max_tokens"));
1666 }
1667
1668 #[test]
1669 fn parse_anthropic_message_start_without_usage_is_empty() {
1670 let data = r#"{"type":"message_start","message":{"id":"msg_1"}}"#;
1672 assert!(parse_anthropic_sse_line("message_start", data).is_empty());
1673 }
1674
1675 #[test]
1676 fn accumulator_tracks_usage_across_anthropic_stream() {
1677 let mut acc = StreamAccumulator::default();
1680 for event in parse_anthropic_sse_line(
1681 "message_start",
1682 r#"{"message":{"usage":{"input_tokens":245,"output_tokens":1}}}"#,
1683 ) {
1684 acc.push(&event);
1685 }
1686 for event in parse_anthropic_sse_line(
1687 "content_block_start",
1688 r#"{"index":0,"content_block":{"type":"text","text":""}}"#,
1689 ) {
1690 acc.push(&event);
1691 }
1692 for (chunk, _) in [
1693 (r#"{"delta":{"type":"text_delta","text":"Hello"}}"#, ()),
1694 (r#"{"delta":{"type":"text_delta","text":", "}}"#, ()),
1695 (r#"{"delta":{"type":"text_delta","text":"world"}}"#, ()),
1696 ] {
1697 for event in parse_anthropic_sse_line("content_block_delta", chunk) {
1698 acc.push(&event);
1699 }
1700 }
1701 for event in parse_anthropic_sse_line("message_delta", r#"{"usage":{"output_tokens":87}}"#)
1702 {
1703 acc.push(&event);
1704 }
1705
1706 let (text, tools, usage, _stop) = acc.finish_with_usage();
1707 assert_eq!(text, "Hello, world");
1708 assert!(tools.is_empty());
1709 let usage = usage.expect("provider reported usage; must surface");
1710 assert_eq!(usage.prompt_tokens, 245);
1711 assert_eq!(usage.completion_tokens, 87);
1713 assert_eq!(usage.total_tokens, 332);
1714 }
1715
1716 #[test]
1717 fn parse_openai_final_chunk_emits_usage() {
1718 let line = r#"data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[],"usage":{"prompt_tokens":245,"completion_tokens":87,"total_tokens":332}}"#;
1721 let events = parse_openai_sse_line(line);
1722 assert_eq!(events.len(), 1);
1723 match &events[0] {
1724 StreamEvent::Usage {
1725 input_tokens,
1726 output_tokens,
1727 ..
1728 } => {
1729 assert_eq!(*input_tokens, 245);
1730 assert_eq!(*output_tokens, 87);
1731 }
1732 other => panic!("expected Usage, got {:?}", other),
1733 }
1734 }
1735
1736 #[test]
1737 fn accumulator_tracks_usage_across_openai_stream() {
1738 let mut acc = StreamAccumulator::default();
1741 for line in [
1742 r#"data: {"choices":[{"delta":{"content":"Hello"}}]}"#,
1743 r#"data: {"choices":[{"delta":{"content":", "}}]}"#,
1744 r#"data: {"choices":[{"delta":{"content":"world"}}]}"#,
1745 r#"data: {"id":"chatcmpl-1","choices":[],"usage":{"prompt_tokens":245,"completion_tokens":87}}"#,
1746 ] {
1747 for event in parse_openai_sse_line(line) {
1748 acc.push(&event);
1749 }
1750 }
1751
1752 let (text, tools, usage, _stop) = acc.finish_with_usage();
1753 assert_eq!(text, "Hello, world");
1754 assert!(tools.is_empty());
1755 let usage = usage.expect("provider reported usage; must surface");
1756 assert_eq!(usage.prompt_tokens, 245);
1757 assert_eq!(usage.completion_tokens, 87);
1758 assert_eq!(usage.total_tokens, 332);
1759 }
1760
1761 #[test]
1762 fn accumulator_returns_no_usage_when_provider_silent() {
1763 let mut acc = StreamAccumulator::default();
1767 acc.push(&StreamEvent::TextDelta("hi".into()));
1768 let (_, _, usage, _stop) = acc.finish_with_usage();
1769 assert!(usage.is_none());
1770 }
1771
1772 #[test]
1773 fn anthropic_stream_decodes_cache_tokens_from_message_start() {
1774 let mut acc = StreamAccumulator::default();
1778 let start = r#"{"message":{"usage":{"input_tokens":50,"output_tokens":1,"cache_read_input_tokens":4000,"cache_creation_input_tokens":600}}}"#;
1779 for e in parse_anthropic_sse_line("message_start", start) {
1780 acc.push(&e);
1781 }
1782 let delta = r#"{"delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":87}}"#;
1783 for e in parse_anthropic_sse_line("message_delta", delta) {
1784 acc.push(&e);
1785 }
1786 let (_t, _c, usage, stop) = acc.finish_with_usage();
1787 let u = usage.expect("usage surfaced");
1788 assert_eq!(u.prompt_tokens, 50, "uncached prefix");
1789 assert_eq!(u.completion_tokens, 87, "final output from message_delta");
1790 assert_eq!(u.cache_read_input_tokens, 4000);
1791 assert_eq!(u.cache_creation_input_tokens, 600);
1792 assert_eq!(stop.as_deref(), Some("end_turn"));
1793 }
1794
1795 #[test]
1796 fn openai_stream_normalizes_cached_tokens_out_of_prompt() {
1797 let mut acc = StreamAccumulator::default();
1801 let chunk = r#"data: {"choices":[],"usage":{"prompt_tokens":1000,"completion_tokens":40,"prompt_tokens_details":{"cached_tokens":800}}}"#;
1802 for e in parse_openai_sse_line(chunk) {
1803 acc.push(&e);
1804 }
1805 let (_t, _c, usage, _s) = acc.finish_with_usage();
1806 let u = usage.expect("usage surfaced");
1807 assert_eq!(u.prompt_tokens, 200, "uncached = 1000 - 800");
1808 assert_eq!(u.cache_read_input_tokens, 800);
1809 assert_eq!(
1810 u.cache_creation_input_tokens, 0,
1811 "OpenAI has no write bucket"
1812 );
1813 }
1814
1815 #[test]
1816 fn responses_failure_is_a_terminal_safe_error_event() {
1817 let events = parse_openai_responses_sse_line(
1818 "response.failed",
1819 r#"{"response":{"error":{"message":"managed model unavailable","stack":"secret"}}}"#,
1820 );
1821 assert!(matches!(
1822 events.as_slice(),
1823 [StreamEvent::Error(message)] if message == "managed model unavailable"
1824 ));
1825 assert!(!format!("{events:?}").contains("secret"));
1826 }
1827
1828 #[test]
1829 fn responses_reasoning_item_done_is_retained_verbatim() {
1830 let data = r#"{"output_index":0,"item":{"type":"reasoning","id":"rs_1","status":"completed","summary":[{"type":"summary_text","text":"safe summary"}],"encrypted_content":"opaque-ciphertext"}}"#;
1831 let events = parse_openai_responses_sse_line("response.output_item.done", data);
1832 let expected = serde_json::json!({
1833 "type": "reasoning",
1834 "id": "rs_1",
1835 "status": "completed",
1836 "summary": [{"type": "summary_text", "text": "safe summary"}],
1837 "encrypted_content": "opaque-ciphertext",
1838 });
1839 assert!(matches!(
1840 events.as_slice(),
1841 [StreamEvent::ProviderOutputItem(item)] if item == &expected
1842 ));
1843 let mut accumulator = StreamAccumulator::default();
1844 accumulator.push(&events[0]);
1845 let (_, _, _, _, items) = accumulator.finish_with_provider_output_items();
1846 assert_eq!(items, vec![expected]);
1847 }
1848
1849 #[test]
1850 fn responses_incomplete_is_terminal_failure_not_success() {
1851 let events = parse_openai_responses_sse_line(
1852 "response.incomplete",
1853 r#"{"response":{"status":"incomplete","incomplete_details":{"reason":"max_output_tokens"},"usage":{"input_tokens":17,"output_tokens":9}}}"#,
1854 );
1855 assert!(
1856 events
1857 .iter()
1858 .any(|event| matches!(event, StreamEvent::Error(_))),
1859 "response.incomplete must emit a terminal error"
1860 );
1861 }
1862
1863 fn run_stop_filter(stops: &[&str], events: Vec<StreamEvent>) -> Vec<StreamEvent> {
1864 let mut filter = StopSequenceFilter::new(stops.iter().map(|s| s.to_string()).collect());
1865 let mut out = Vec::new();
1866 for event in events {
1867 out.extend(filter.push(event));
1868 }
1869 out.extend(filter.finish());
1870 out
1871 }
1872
1873 fn streamed_text(events: &[StreamEvent]) -> String {
1874 events
1875 .iter()
1876 .filter_map(|event| match event {
1877 StreamEvent::TextDelta(delta) => Some(delta.as_str()),
1878 _ => None,
1879 })
1880 .collect()
1881 }
1882
1883 #[test]
1884 fn stop_filter_catches_a_sequence_split_across_three_deltas() {
1885 let out = run_stop_filter(
1886 &["STOPHERE"],
1887 vec![
1888 StreamEvent::TextDelta("Paris S".into()),
1889 StreamEvent::TextDelta("TOPH".into()),
1890 StreamEvent::Usage {
1891 input_tokens: 3,
1892 output_tokens: 2,
1893 cache_read_input_tokens: 0,
1894 cache_creation_input_tokens: 0,
1895 },
1896 StreamEvent::TextDelta("ERE indeed".into()),
1897 StreamEvent::TextDelta(" and more".into()),
1898 StreamEvent::StopReason("length".into()),
1899 StreamEvent::Done {
1900 text: "Paris STOPHERE indeed and more".into(),
1901 tool_calls: Vec::new(),
1902 },
1903 ],
1904 );
1905 assert_eq!(streamed_text(&out), "Paris ");
1906 let reasons: Vec<&str> = out
1907 .iter()
1908 .filter_map(|event| match event {
1909 StreamEvent::StopReason(reason) => Some(reason.as_str()),
1910 _ => None,
1911 })
1912 .collect();
1913 assert_eq!(reasons, ["stop_sequence"]);
1914 assert!(matches!(
1915 out.last(),
1916 Some(StreamEvent::Done { text, .. }) if text == "Paris "
1917 ));
1918 assert!(out
1920 .iter()
1921 .any(|event| matches!(event, StreamEvent::Usage { .. })));
1922 }
1923
1924 #[test]
1927 fn stop_filter_releases_a_false_alarm_tail_unchanged() {
1928 let out = run_stop_filter(
1929 &["STOPHERE"],
1930 vec![
1931 StreamEvent::TextDelta("almost STO".into()),
1932 StreamEvent::TextDelta("P but not quite".into()),
1933 StreamEvent::StopReason("stop".into()),
1934 StreamEvent::Done {
1935 text: "almost STOP but not quite".into(),
1936 tool_calls: Vec::new(),
1937 },
1938 ],
1939 );
1940 assert_eq!(streamed_text(&out), "almost STOP but not quite");
1941 assert!(out
1942 .iter()
1943 .any(|event| matches!(event, StreamEvent::StopReason(reason) if reason == "stop")));
1944 assert!(matches!(
1945 out.last(),
1946 Some(StreamEvent::Done { text, .. }) if text == "almost STOP but not quite"
1947 ));
1948 }
1949
1950 #[test]
1954 fn stop_filter_reports_one_reason_when_only_done_carries_the_stop() {
1955 let out = run_stop_filter(
1956 &["STOPHERE"],
1957 vec![
1958 StreamEvent::StopReason("length".into()),
1959 StreamEvent::Done {
1960 text: "Paris STOPHERE indeed".into(),
1961 tool_calls: Vec::new(),
1962 },
1963 ],
1964 );
1965 let reasons: Vec<&str> = out
1966 .iter()
1967 .filter_map(|event| match event {
1968 StreamEvent::StopReason(reason) => Some(reason.as_str()),
1969 _ => None,
1970 })
1971 .collect();
1972 assert_eq!(reasons, ["stop_sequence"]);
1973 assert!(matches!(
1974 out.last(),
1975 Some(StreamEvent::Done { text, .. }) if text == "Paris "
1976 ));
1977 }
1978
1979 #[test]
1980 fn stop_filter_flushes_a_held_tail_when_the_source_closes() {
1981 let out = run_stop_filter(
1982 &["STOPHERE"],
1983 vec![StreamEvent::TextDelta("ends ST".into())],
1984 );
1985 assert_eq!(streamed_text(&out), "ends ST");
1986 }
1987
1988 #[test]
1989 fn stop_filter_holds_back_only_at_char_boundaries() {
1990 let out = run_stop_filter(
1991 &["\u{e9}\u{e9}"],
1992 vec![
1993 StreamEvent::TextDelta("caf\u{e9}".into()),
1994 StreamEvent::TextDelta("\u{e9} au lait".into()),
1995 ],
1996 );
1997 assert_eq!(streamed_text(&out), "caf");
1998 }
1999
2000 #[test]
2001 fn apply_stop_backstop_cuts_text_and_reports_stop_sequence() {
2002 use crate::tasks::generate::apply_stop_backstop;
2003 let stops = vec!["STOPHERE".to_string()];
2004 assert_eq!(
2005 apply_stop_backstop(
2006 "Paris STOPHERE indeed".into(),
2007 Some("length".into()),
2008 &stops,
2009 0
2010 ),
2011 ("Paris ".to_string(), Some("stop_sequence".to_string()))
2012 );
2013 assert_eq!(
2015 apply_stop_backstop("Paris".into(), Some("stop".into()), &stops, 0),
2016 ("Paris".to_string(), Some("stop".to_string()))
2017 );
2018 assert_eq!(
2020 apply_stop_backstop("a STOPHERE b".into(), None, &[], 0),
2021 ("a STOPHERE b".to_string(), None)
2022 );
2023 }
2024}