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
133pub fn parse_openai_responses_sse_line(event_type: &str, data: &str) -> Vec<StreamEvent> {
141 let json: serde_json::Value = match serde_json::from_str(data) {
142 Ok(v) => v,
143 Err(_) => return Vec::new(),
144 };
145 let mut events = Vec::new();
146 match event_type {
147 "response.output_text.delta" => {
148 if let Some(d) = json.get("delta").and_then(|d| d.as_str()) {
149 if !d.is_empty() {
150 events.push(StreamEvent::TextDelta(d.to_string()));
151 }
152 }
153 }
154 "response.output_item.added" | "response.output_item.done" => {
155 if let Some(item) = json.get("item") {
156 if item.get("type").and_then(|t| t.as_str()) == Some("function_call") {
157 let name = item
158 .get("name")
159 .and_then(|n| n.as_str())
160 .unwrap_or("")
161 .to_string();
162 let id = item
163 .get("call_id")
164 .or_else(|| item.get("id"))
165 .and_then(|i| i.as_str())
166 .map(|s| s.to_string());
167 let index = json
168 .get("output_index")
169 .and_then(|v| v.as_u64())
170 .unwrap_or(0) as usize;
171 if !name.is_empty() && event_type == "response.output_item.added" {
175 events.push(StreamEvent::ToolCallStart { name, index, id });
176 }
177 if let Some(arguments) = item
178 .get("arguments")
179 .and_then(|value| value.as_str())
180 .filter(|arguments| !arguments.is_empty())
181 {
182 events.push(StreamEvent::ToolCallArguments {
183 index,
184 arguments: arguments.to_string(),
185 authoritative: event_type == "response.output_item.done",
186 });
187 }
188 } else if event_type == "response.output_item.done"
189 && item.get("type").and_then(|value| value.as_str()) == Some("reasoning")
190 {
191 events.push(StreamEvent::ProviderOutputItem(item.clone()));
192 }
193 }
194 }
195 "response.function_call_arguments.delta" => {
196 if let Some(d) = json.get("delta").and_then(|d| d.as_str()) {
197 let index = json
198 .get("output_index")
199 .and_then(|v| v.as_u64())
200 .unwrap_or(0) as usize;
201 events.push(StreamEvent::ToolCallDelta {
202 index,
203 arguments_delta: d.to_string(),
204 });
205 }
206 }
207 "response.function_call_arguments.done" => {
208 if let Some(arguments) = json.get("arguments").and_then(|value| value.as_str()) {
209 let index = json
210 .get("output_index")
211 .and_then(|value| value.as_u64())
212 .unwrap_or(0) as usize;
213 events.push(StreamEvent::ToolCallArguments {
214 index,
215 arguments: arguments.to_string(),
216 authoritative: true,
217 });
218 }
219 }
220 "response.completed" | "response.incomplete" => {
221 if let Some(resp) = json.get("response") {
222 if let Some(u) = resp.get("usage") {
223 let input_total = u.get("input_tokens").and_then(|v| v.as_u64()).unwrap_or(0);
227 let cached = u
228 .get("input_tokens_details")
229 .and_then(|d| d.get("cached_tokens"))
230 .and_then(|v| v.as_u64())
231 .unwrap_or(0)
232 .min(input_total);
233 events.push(StreamEvent::Usage {
234 input_tokens: input_total - cached,
235 output_tokens: u.get("output_tokens").and_then(|v| v.as_u64()).unwrap_or(0),
236 cache_read_input_tokens: cached,
237 cache_creation_input_tokens: 0,
238 });
239 }
240 if event_type == "response.incomplete" {
241 let reason = resp
242 .pointer("/incomplete_details/reason")
243 .and_then(|r| r.as_str())
244 .unwrap_or("incomplete");
245 events.push(StreamEvent::StopReason(reason.to_string()));
246 events.push(StreamEvent::Error(format!(
247 "managed inference incomplete: {reason}"
248 )));
249 } else {
250 events.push(StreamEvent::Done {
255 text: String::new(),
256 tool_calls: Vec::new(),
257 });
258 }
259 }
260 }
261 "error" | "response.failed" => {
262 let pick = |field: &str| {
263 json.pointer(&format!("/error/{field}"))
264 .or_else(|| json.pointer(&format!("/response/error/{field}")))
265 .and_then(|value| value.as_str())
266 .map(str::trim)
267 .filter(|value| !value.is_empty())
268 .map(str::to_string)
269 };
270 let message = pick("message");
271 let kind = pick("type");
285 let code = pick("code");
286 let mut detail = message.unwrap_or_else(|| "managed inference failed".to_string());
287 let tags: Vec<String> = [("type", kind), ("code", code)]
288 .into_iter()
289 .filter_map(|(label, value)| value.map(|v| format!("{label}={v}")))
290 .collect();
291 if !tags.is_empty() {
292 detail.push_str(&format!(" ({})", tags.join(", ")));
293 }
294 events.push(StreamEvent::Error(detail));
298 }
299 _ => {}
300 }
301 events
302}
303
304pub fn error_tags(detail: &str) -> (Option<&str>, Option<&str>) {
316 let Some(open) = detail.rfind(" (") else {
317 return (None, None);
318 };
319 let Some(close) = detail[open..].rfind(')') else {
320 return (None, None);
321 };
322 let mut kind = None;
323 let mut code = None;
324 for part in detail[open + 2..open + close].split(", ") {
325 if let Some(v) = part.strip_prefix("type=") {
326 kind = Some(v);
327 } else if let Some(v) = part.strip_prefix("code=") {
328 code = Some(v);
329 }
330 }
331 (kind, code)
332}
333
334pub fn content_refusal_tags(detail: &str) -> Option<(Option<String>, Option<String>)> {
358 let (kind, code) = error_tags(detail);
359 let refused = |v: &str| {
360 let v = v.to_ascii_lowercase();
361 v == "refusal"
363 || v.contains("content_policy")
364 || v.contains("content_filter")
365 || v.contains("moderation")
366 || v.contains("safety")
367 };
368 (code.is_some_and(refused) || kind.is_some_and(refused))
369 .then(|| (kind.map(str::to_string), code.map(str::to_string)))
370}
371
372pub(crate) fn google_output_tokens(usage: &serde_json::Value) -> u64 {
381 let count = |field: &str| usage.get(field).and_then(|v| v.as_u64()).unwrap_or(0);
382 count("candidatesTokenCount") + count("thoughtsTokenCount")
383}
384
385pub fn parse_google_sse_line(data: &str) -> Vec<StreamEvent> {
393 let json: serde_json::Value = match serde_json::from_str(data) {
394 Ok(v) => v,
395 Err(_) => return Vec::new(),
396 };
397 let mut events = Vec::new();
398 if let Some(parts) = json
399 .pointer("/candidates/0/content/parts")
400 .and_then(|p| p.as_array())
401 {
402 for (i, part) in parts.iter().enumerate() {
403 if let Some(t) = part.get("text").and_then(|t| t.as_str()) {
404 if !t.is_empty() {
405 events.push(StreamEvent::TextDelta(t.to_string()));
406 }
407 }
408 if let Some(fc) = part.get("functionCall") {
409 let name = fc
410 .get("name")
411 .and_then(|n| n.as_str())
412 .unwrap_or("")
413 .to_string();
414 if !name.is_empty() {
415 let args = fc
416 .get("args")
417 .map(|a| a.to_string())
418 .unwrap_or_else(|| "{}".to_string());
419 events.push(StreamEvent::ToolCallStart {
420 name,
421 index: i,
422 id: None,
423 });
424 events.push(StreamEvent::ToolCallDelta {
425 index: i,
426 arguments_delta: args,
427 });
428 }
429 }
430 }
431 }
432 if let Some(u) = json.get("usageMetadata") {
433 events.push(StreamEvent::Usage {
434 input_tokens: u
435 .get("promptTokenCount")
436 .and_then(|v| v.as_u64())
437 .unwrap_or(0),
438 output_tokens: google_output_tokens(u),
439 cache_read_input_tokens: 0,
441 cache_creation_input_tokens: 0,
442 });
443 }
444 if let Some(fr) = json
445 .pointer("/candidates/0/finishReason")
446 .and_then(|r| r.as_str())
447 {
448 events.push(StreamEvent::StopReason(fr.to_string()));
449 }
450 events
451}
452
453pub(crate) fn fold_openai_chat_sse(body: &str) -> Result<String, crate::InferenceError> {
466 if body.trim_start().starts_with('{') {
468 return Ok(body.to_string());
469 }
470 let mut chunks = 0usize;
471 let mut done = false;
472 let mut content = String::new();
473 let mut reasoning = String::new();
474 let mut tools: Vec<(u64, String, String, String)> = Vec::new();
476 let mut finish_reason = serde_json::Value::Null;
477 let mut usage = serde_json::Value::Null;
478 for line in body.lines() {
479 let Some(data) = line.strip_prefix("data:").map(str::trim) else {
480 continue;
481 };
482 if data == "[DONE]" {
483 done = true;
484 continue;
485 }
486 if data.is_empty() {
487 continue;
488 }
489 let chunk: serde_json::Value = serde_json::from_str(data).map_err(|e| {
490 crate::InferenceError::InferenceFailed(format!("parse stream chunk: {e}"))
491 })?;
492 chunks += 1;
493 if let Some(error) = chunk.get("error") {
494 let message = error
495 .get("message")
496 .and_then(|m| m.as_str())
497 .unwrap_or("the server reported an error mid-stream")
498 .to_string();
499 let status = error.get("code").and_then(|c| c.as_u64()).map(|c| c as u16);
503 return Err(match status {
504 Some(code) if code == 429 || code >= 500 => crate::InferenceError::Transient {
505 status: Some(code),
506 message,
507 },
508 _ => crate::InferenceError::InferenceFailed(message),
509 });
510 }
511 if let Some(u) = chunk.get("usage").filter(|u| !u.is_null()) {
512 usage = u.clone();
513 }
514 let Some(choice) = chunk.pointer("/choices/0") else {
515 continue;
516 };
517 if let Some(fr) = choice.get("finish_reason").filter(|f| !f.is_null()) {
518 finish_reason = fr.clone();
519 }
520 let Some(delta) = choice.get("delta") else {
521 continue;
522 };
523 if let Some(t) = delta.get("content").and_then(|c| c.as_str()) {
524 content.push_str(t);
525 }
526 if let Some(t) = delta
527 .get("reasoning_content")
528 .or_else(|| delta.get("reasoning"))
529 .and_then(|c| c.as_str())
530 {
531 reasoning.push_str(t);
532 }
533 for call in delta
534 .get("tool_calls")
535 .and_then(|t| t.as_array())
536 .into_iter()
537 .flatten()
538 {
539 let index = call.get("index").and_then(|i| i.as_u64()).unwrap_or(0);
540 let slot = match tools.iter().position(|t| t.0 == index) {
541 Some(i) => &mut tools[i],
542 None => {
543 tools.push((index, String::new(), String::new(), String::new()));
544 tools.last_mut().expect("just pushed")
545 }
546 };
547 if let Some(id) = call.get("id").and_then(|v| v.as_str()) {
548 slot.1 = id.to_string();
549 }
550 if let Some(name) = call.pointer("/function/name").and_then(|v| v.as_str()) {
551 slot.2.push_str(name);
552 }
553 if let Some(args) = call.pointer("/function/arguments").and_then(|v| v.as_str()) {
554 slot.3.push_str(args);
555 }
556 }
557 }
558 if chunks == 0 || (!done && finish_reason.is_null()) {
561 return Err(crate::InferenceError::InferenceFailed(
562 "the stream ended before the server finished the response".to_string(),
563 ));
564 }
565 let mut message = serde_json::json!({ "role": "assistant", "content": content });
566 if !reasoning.is_empty() {
567 message["reasoning_content"] = serde_json::Value::String(reasoning);
568 }
569 if !tools.is_empty() {
570 message["tool_calls"] = tools
571 .into_iter()
572 .map(|(_, id, name, arguments)| {
573 serde_json::json!({
574 "id": id,
575 "type": "function",
576 "function": { "name": name, "arguments": arguments },
577 })
578 })
579 .collect();
580 }
581 let mut response = serde_json::json!({
582 "choices": [{ "index": 0, "message": message, "finish_reason": finish_reason }],
583 });
584 if !usage.is_null() {
586 response["usage"] = usage;
587 }
588 Ok(response.to_string())
589}
590
591pub fn parse_openai_sse_line(line: &str) -> Vec<StreamEvent> {
594 let data = match line.strip_prefix("data: ") {
595 Some(d) => d,
596 None => return Vec::new(),
597 };
598 if data == "[DONE]" {
599 return Vec::new();
600 }
601
602 let json: serde_json::Value = match serde_json::from_str(data) {
603 Ok(v) => v,
604 Err(_) => return Vec::new(),
605 };
606
607 let mut events = Vec::new();
608
609 if let Some(reason) = json
613 .get("choices")
614 .and_then(|c| c.as_array())
615 .and_then(|c| c.first())
616 .and_then(|c| c.get("finish_reason"))
617 .and_then(|r| r.as_str())
618 {
619 if !reason.is_empty() {
620 events.push(StreamEvent::StopReason(reason.to_string()));
621 }
622 }
623
624 if let Some(delta) = json
627 .get("choices")
628 .and_then(|c| c.as_array())
629 .and_then(|c| c.first())
630 .and_then(|c| c.get("delta"))
631 {
632 if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
633 if !content.is_empty() {
634 events.push(StreamEvent::TextDelta(content.to_string()));
635 }
636 }
637
638 if let Some(tool_calls) = delta.get("tool_calls").and_then(|t| t.as_array()) {
640 for tc in tool_calls {
641 let index = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
642 if let Some(function) = tc.get("function") {
643 if let Some(name) = function.get("name").and_then(|n| n.as_str()) {
644 let id = tc.get("id").and_then(|i| i.as_str()).map(|s| s.to_string());
645 events.push(StreamEvent::ToolCallStart {
646 name: name.to_string(),
647 index,
648 id,
649 });
650 }
651 if let Some(args) = function.get("arguments").and_then(|a| a.as_str()) {
652 if !args.is_empty() {
653 events.push(StreamEvent::ToolCallDelta {
654 index,
655 arguments_delta: args.to_string(),
656 });
657 }
658 }
659 }
660 }
661 }
662 }
663
664 if let Some(usage) = json.get("usage") {
668 let input = usage
669 .get("prompt_tokens")
670 .and_then(|n| n.as_u64())
671 .unwrap_or(0);
672 let output = usage
673 .get("completion_tokens")
674 .and_then(|n| n.as_u64())
675 .unwrap_or(0);
676 let cached = usage
679 .get("prompt_tokens_details")
680 .and_then(|d| d.get("cached_tokens"))
681 .and_then(|n| n.as_u64())
682 .unwrap_or(0)
683 .min(input);
684 if input != 0 || output != 0 {
685 events.push(StreamEvent::Usage {
686 input_tokens: input - cached,
687 output_tokens: output,
688 cache_read_input_tokens: cached,
689 cache_creation_input_tokens: 0,
690 });
691 }
692 }
693
694 events
695}
696
697pub fn parse_anthropic_sse_line(event_type: &str, data: &str) -> Vec<StreamEvent> {
699 match event_type {
700 "content_block_delta" => {
701 let json: serde_json::Value = match serde_json::from_str(data) {
702 Ok(v) => v,
703 Err(_) => return Vec::new(),
704 };
705 let delta = match json.get("delta") {
706 Some(d) => d,
707 None => return Vec::new(),
708 };
709 let delta_type = match delta.get("type").and_then(|t| t.as_str()) {
710 Some(t) => t,
711 None => return Vec::new(),
712 };
713
714 match delta_type {
715 "text_delta" => match delta.get("text").and_then(|t| t.as_str()) {
716 Some(text) => vec![StreamEvent::TextDelta(text.to_string())],
717 None => Vec::new(),
718 },
719 "input_json_delta" => match delta.get("partial_json").and_then(|p| p.as_str()) {
720 Some(partial) => {
721 let index =
722 json.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
723 vec![StreamEvent::ToolCallDelta {
724 index,
725 arguments_delta: partial.to_string(),
726 }]
727 }
728 None => Vec::new(),
729 },
730 _ => Vec::new(),
731 }
732 }
733 "content_block_start" => {
734 let json: serde_json::Value = match serde_json::from_str(data) {
735 Ok(v) => v,
736 Err(_) => return Vec::new(),
737 };
738 let block = match json.get("content_block") {
739 Some(b) => b,
740 None => return Vec::new(),
741 };
742 if block.get("type").and_then(|t| t.as_str()) == Some("tool_use") {
743 if let Some(name) = block.get("name").and_then(|n| n.as_str()) {
744 let index = json.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
745 let id = block
746 .get("id")
747 .and_then(|i| i.as_str())
748 .map(|s| s.to_string());
749 return vec![StreamEvent::ToolCallStart {
750 name: name.to_string(),
751 index,
752 id,
753 }];
754 }
755 }
756 Vec::new()
757 }
758 "message_start" => {
762 let json: serde_json::Value = match serde_json::from_str(data) {
763 Ok(v) => v,
764 Err(_) => return Vec::new(),
765 };
766 let Some(usage) = json.pointer("/message/usage") else {
767 return Vec::new();
768 };
769 let input = usage
770 .get("input_tokens")
771 .and_then(|n| n.as_u64())
772 .unwrap_or(0);
773 let output = usage
774 .get("output_tokens")
775 .and_then(|n| n.as_u64())
776 .unwrap_or(0);
777 let cache_read = usage
780 .get("cache_read_input_tokens")
781 .and_then(|n| n.as_u64())
782 .unwrap_or(0);
783 let cache_creation = usage
784 .get("cache_creation_input_tokens")
785 .and_then(|n| n.as_u64())
786 .unwrap_or(0);
787 if input == 0 && output == 0 && cache_read == 0 && cache_creation == 0 {
788 return Vec::new();
789 }
790 vec![StreamEvent::Usage {
791 input_tokens: input,
792 output_tokens: output,
793 cache_read_input_tokens: cache_read,
794 cache_creation_input_tokens: cache_creation,
795 }]
796 }
797 "message_delta" => {
801 let json: serde_json::Value = match serde_json::from_str(data) {
802 Ok(v) => v,
803 Err(_) => return Vec::new(),
804 };
805 let mut events = Vec::new();
806 if let Some(reason) = json.pointer("/delta/stop_reason").and_then(|r| r.as_str()) {
809 if !reason.is_empty() {
810 events.push(StreamEvent::StopReason(reason.to_string()));
811 }
812 }
813 if let Some(usage) = json.get("usage") {
814 let input = usage
815 .get("input_tokens")
816 .and_then(|n| n.as_u64())
817 .unwrap_or(0);
818 let output = usage
819 .get("output_tokens")
820 .and_then(|n| n.as_u64())
821 .unwrap_or(0);
822 if input != 0 || output != 0 {
823 events.push(StreamEvent::Usage {
824 input_tokens: input,
825 output_tokens: output,
826 cache_read_input_tokens: 0,
829 cache_creation_input_tokens: 0,
830 });
831 }
832 }
833 events
834 }
835 _ => Vec::new(),
836 }
837}
838
839pub fn enforce_stop_sequences(
860 mut source: tokio::sync::mpsc::Receiver<StreamEvent>,
861 stops: Vec<String>,
862) -> tokio::sync::mpsc::Receiver<StreamEvent> {
863 let stops: Vec<String> = stops.into_iter().filter(|stop| !stop.is_empty()).collect();
864 if stops.is_empty() {
865 return source;
866 }
867 let (tx, rx) = tokio::sync::mpsc::channel(64);
868 tokio::spawn(async move {
869 let mut filter = StopSequenceFilter::new(stops);
870 while let Some(event) = source.recv().await {
871 for out in filter.push(event) {
872 if tx.send(out).await.is_err() {
873 return;
874 }
875 }
876 }
877 for out in filter.finish() {
878 if tx.send(out).await.is_err() {
879 return;
880 }
881 }
882 });
883 rx
884}
885
886struct StopSequenceFilter {
889 stops: Vec<String>,
890 held: String,
893 stopped: bool,
895 reported_stop: bool,
897 provider_reason: Option<String>,
900}
901
902impl StopSequenceFilter {
903 fn new(stops: Vec<String>) -> Self {
904 Self {
905 stops,
906 held: String::new(),
907 stopped: false,
908 reported_stop: false,
909 provider_reason: None,
910 }
911 }
912
913 fn ambiguous_tail(&self) -> usize {
916 self.held
917 .char_indices()
918 .map(|(index, _)| index)
919 .find(|&index| {
920 let tail = &self.held[index..];
921 self.stops
922 .iter()
923 .any(|stop| stop.len() > tail.len() && stop.starts_with(tail))
924 })
925 .map_or(0, |index| self.held.len() - index)
926 }
927
928 fn take_held(&mut self) -> Vec<StreamEvent> {
929 if self.held.is_empty() {
930 Vec::new()
931 } else {
932 vec![StreamEvent::TextDelta(std::mem::take(&mut self.held))]
933 }
934 }
935
936 fn stop_reason(&mut self) -> StreamEvent {
937 self.reported_stop = true;
938 StreamEvent::StopReason(crate::tasks::generate::STOP_SEQUENCE_STOP_REASON.to_string())
939 }
940
941 fn push(&mut self, event: StreamEvent) -> Vec<StreamEvent> {
942 match event {
943 StreamEvent::TextDelta(delta) => {
944 if self.stopped {
945 return Vec::new();
946 }
947 self.held.push_str(&delta);
948 let earliest = self
949 .stops
950 .iter()
951 .filter_map(|stop| self.held.find(stop.as_str()))
952 .min();
953 if let Some(cut) = earliest {
954 self.held.truncate(cut);
955 self.stopped = true;
956 return self.take_held();
957 }
958 let emit = self.held.len() - self.ambiguous_tail();
959 if emit == 0 {
960 return Vec::new();
961 }
962 let rest = self.held.split_off(emit);
963 vec![StreamEvent::TextDelta(std::mem::replace(
964 &mut self.held,
965 rest,
966 ))]
967 }
968 StreamEvent::StopReason(reason) => {
969 self.provider_reason = Some(reason);
972 Vec::new()
973 }
974 StreamEvent::Done { text, tool_calls } => {
975 let mut out = self.take_held();
976 let cut = crate::tasks::generate::truncate_at_stop(&text, &self.stops);
977 if cut.len() < text.len() {
978 self.stopped = true;
979 }
980 out.extend(self.reason_event());
981 if !tool_calls.is_empty() {
982 tracing::warn!(
983 tool_calls = tool_calls.len(),
984 "stop sequences apply to streamed text only; tool-call arguments \
985 in this stream were not checked against them"
986 );
987 }
988 out.push(StreamEvent::Done {
989 text: cut,
990 tool_calls,
991 });
992 out
993 }
994 StreamEvent::Error(message) => {
995 let mut out = self.take_held();
997 out.extend(self.reason_event());
998 out.push(StreamEvent::Error(message));
999 out
1000 }
1001 other @ (StreamEvent::ToolCallStart { .. }
1005 | StreamEvent::ToolCallDelta { .. }
1006 | StreamEvent::ToolCallArguments { .. }
1007 | StreamEvent::Usage { .. }
1008 | StreamEvent::ProviderCost(_)
1009 | StreamEvent::ProviderOutputItem(_)
1010 | StreamEvent::Attribution(_)
1011 | StreamEvent::Substituted(_)) => vec![other],
1012 }
1013 }
1014
1015 fn reason_event(&mut self) -> Option<StreamEvent> {
1018 if self.reported_stop {
1019 return None;
1020 }
1021 if self.stopped {
1022 self.provider_reason = None;
1023 return Some(self.stop_reason());
1024 }
1025 self.provider_reason.take().map(|reason| {
1026 self.reported_stop = true;
1027 StreamEvent::StopReason(reason)
1028 })
1029 }
1030
1031 fn finish(&mut self) -> Vec<StreamEvent> {
1033 let mut out = self.take_held();
1034 out.extend(self.reason_event());
1035 out
1036 }
1037}
1038
1039#[derive(Default)]
1041pub struct StreamAccumulator {
1042 pub text: String,
1043 tool_names: HashMap<usize, String>,
1044 tool_args: HashMap<usize, String>,
1045 tool_args_fallback: HashMap<usize, String>,
1046 tool_args_authoritative: HashMap<usize, String>,
1047 tool_ids: HashMap<usize, String>,
1048 input_tokens: u64,
1053 output_tokens: u64,
1057 cache_read_input_tokens: u64,
1059 cache_creation_input_tokens: u64,
1061 saw_usage: bool,
1066 stop_reason: Option<String>,
1069 provider_output_items: Vec<serde_json::Value>,
1071 reported_cost_usd: Option<f64>,
1073}
1074
1075impl StreamAccumulator {
1076 pub fn push(&mut self, event: &StreamEvent) {
1077 match event {
1078 StreamEvent::TextDelta(t) => self.text.push_str(t),
1079 StreamEvent::ToolCallStart { name, index, id } => {
1080 self.tool_names.insert(*index, name.clone());
1081 self.tool_args.entry(*index).or_default();
1082 if let Some(id) = id {
1083 self.tool_ids.insert(*index, id.clone());
1084 }
1085 }
1086 StreamEvent::ToolCallDelta {
1087 index,
1088 arguments_delta,
1089 } => {
1090 self.tool_args
1091 .entry(*index)
1092 .or_default()
1093 .push_str(arguments_delta);
1094 }
1095 StreamEvent::ToolCallArguments {
1096 index,
1097 arguments,
1098 authoritative,
1099 } => {
1100 let destination = if *authoritative {
1101 &mut self.tool_args_authoritative
1102 } else {
1103 &mut self.tool_args_fallback
1104 };
1105 destination.insert(*index, arguments.clone());
1106 }
1107 StreamEvent::Usage {
1108 input_tokens,
1109 output_tokens,
1110 cache_read_input_tokens,
1111 cache_creation_input_tokens,
1112 } => {
1113 self.saw_usage = true;
1114 if *input_tokens > self.input_tokens {
1121 self.input_tokens = *input_tokens;
1122 }
1123 if *output_tokens > self.output_tokens {
1124 self.output_tokens = *output_tokens;
1125 }
1126 if *cache_read_input_tokens > self.cache_read_input_tokens {
1127 self.cache_read_input_tokens = *cache_read_input_tokens;
1128 }
1129 if *cache_creation_input_tokens > self.cache_creation_input_tokens {
1130 self.cache_creation_input_tokens = *cache_creation_input_tokens;
1131 }
1132 }
1133 StreamEvent::StopReason(reason) => {
1134 self.stop_reason = Some(reason.clone());
1135 }
1136 StreamEvent::ProviderOutputItem(item) => {
1137 self.provider_output_items.push(item.clone());
1138 }
1139 StreamEvent::ProviderCost(cost) => {
1140 self.reported_cost_usd = Some(*cost);
1141 }
1142 StreamEvent::Error(_) | StreamEvent::Attribution(_) | StreamEvent::Substituted(_) => {}
1143 StreamEvent::Done { .. } => {}
1144 }
1145 }
1146
1147 pub fn finish(self) -> (String, Vec<ToolCall>) {
1148 let (text, tool_calls, _, _) = self.finish_with_usage();
1149 (text, tool_calls)
1150 }
1151
1152 pub fn finish_with_usage(self) -> (String, Vec<ToolCall>, Option<TokenUsage>, Option<String>) {
1160 let (text, tool_calls, usage, stop_reason, _) = self.finish_with_provider_output_items();
1161 (text, tool_calls, usage, stop_reason)
1162 }
1163
1164 pub fn finish_with_provider_output_items(
1167 self,
1168 ) -> (
1169 String,
1170 Vec<ToolCall>,
1171 Option<TokenUsage>,
1172 Option<String>,
1173 Vec<serde_json::Value>,
1174 ) {
1175 let mut tool_calls = Vec::new();
1176 let mut indices: Vec<usize> = self.tool_names.keys().copied().collect();
1177 indices.sort();
1178
1179 for idx in indices {
1180 let id = self.tool_ids.get(&idx).cloned();
1181 let name = self.tool_names.get(&idx).cloned().unwrap_or_default();
1182 let args_str = self
1183 .tool_args_authoritative
1184 .get(&idx)
1185 .or_else(|| {
1186 self.tool_args
1187 .get(&idx)
1188 .filter(|arguments| !arguments.is_empty())
1189 })
1190 .or_else(|| self.tool_args_fallback.get(&idx))
1191 .cloned()
1192 .unwrap_or_default();
1193 let arguments: HashMap<String, serde_json::Value> = if args_str.trim().is_empty() {
1194 HashMap::new()
1195 } else {
1196 match serde_json::from_str(&args_str) {
1197 Ok(arguments) => arguments,
1198 Err(error) => {
1199 tracing::warn!(
1200 tool = %name,
1201 output_index = idx,
1202 %error,
1203 "tool-call arguments were not valid JSON; returning an empty argument map"
1204 );
1205 HashMap::new()
1206 }
1207 }
1208 };
1209 tool_calls.push(ToolCall {
1210 id,
1211 name,
1212 arguments,
1213 });
1214 }
1215
1216 let usage = if self.saw_usage {
1217 Some(TokenUsage {
1218 prompt_tokens: self.input_tokens,
1219 completion_tokens: self.output_tokens,
1220 total_tokens: self.input_tokens + self.output_tokens,
1221 context_window: 0,
1225 cache_read_input_tokens: self.cache_read_input_tokens,
1228 cache_creation_input_tokens: self.cache_creation_input_tokens,
1229 reported_cost_usd: self.reported_cost_usd,
1230 reasoning_tokens: None,
1232 })
1233 } else {
1234 None
1235 };
1236
1237 let (text, tag_calls) = crate::tasks::generate::parse_tool_calls(&self.text);
1243 let (text, tool_calls) = if tool_calls.is_empty() && !tag_calls.is_empty() {
1244 (text, tag_calls)
1245 } else {
1246 (text, tool_calls)
1248 };
1249
1250 (
1251 text,
1252 tool_calls,
1253 usage,
1254 self.stop_reason,
1255 self.provider_output_items,
1256 )
1257 }
1258}
1259
1260pub fn parse_sse_lines(chunk: &str) -> Vec<(String, String)> {
1263 let mut events = Vec::new();
1264 let mut current_event = String::new();
1265 let mut current_data = String::new();
1266
1267 for line in chunk.lines() {
1268 if let Some(rest) = line.strip_prefix("event: ") {
1269 current_event = rest.to_string();
1270 } else if let Some(rest) = line.strip_prefix("data: ") {
1271 current_data = rest.to_string();
1272 } else if line.is_empty() && !current_data.is_empty() {
1273 events.push((
1274 if current_event.is_empty() {
1275 "message".to_string()
1276 } else {
1277 current_event.clone()
1278 },
1279 current_data.clone(),
1280 ));
1281 current_event.clear();
1282 current_data.clear();
1283 }
1284 }
1285
1286 if !current_data.is_empty() {
1288 events.push((
1289 if current_event.is_empty() {
1290 "message".to_string()
1291 } else {
1292 current_event
1293 },
1294 current_data,
1295 ));
1296 }
1297
1298 events
1299}
1300
1301#[cfg(test)]
1302mod tests {
1303
1304 #[test]
1305 fn a_folded_stream_must_be_a_complete_stream() {
1306 let json = r#"{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}]}"#;
1308 assert_eq!(fold_openai_chat_sse(json).unwrap(), json);
1309 let cut = "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"par\"}}]}\n\n";
1311 assert!(fold_openai_chat_sse(cut).is_err());
1312 assert!(fold_openai_chat_sse("").is_err());
1313 let no_usage = "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n";
1315 let folded: serde_json::Value =
1316 serde_json::from_str(&fold_openai_chat_sse(no_usage).unwrap()).unwrap();
1317 assert!(folded.get("usage").is_none(), "{folded}");
1318 assert_eq!(
1319 folded
1320 .pointer("/choices/0/message/content")
1321 .and_then(|c| c.as_str()),
1322 Some("ok")
1323 );
1324 }
1325
1326 #[test]
1329 fn a_mid_stream_server_error_keeps_its_status() {
1330 let busy = "data: {\"error\":{\"message\":\"Internal server error\",\"code\":500}}\n\n";
1331 assert!(matches!(
1332 fold_openai_chat_sse(busy),
1333 Err(crate::InferenceError::Transient {
1334 status: Some(500),
1335 ..
1336 })
1337 ));
1338 let bad = "data: {\"error\":{\"message\":\"bad request\",\"code\":400}}\n\n";
1339 assert!(matches!(
1340 fold_openai_chat_sse(bad),
1341 Err(crate::InferenceError::InferenceFailed(_))
1342 ));
1343 }
1344
1345 use super::*;
1346
1347 #[test]
1355 fn accumulated_usage_and_stop_reason_survive_finish() {
1356 let mut acc = StreamAccumulator::default();
1357 acc.push(&StreamEvent::TextDelta("hello".into()));
1358 acc.push(&StreamEvent::Usage {
1359 input_tokens: 28,
1360 output_tokens: 5,
1361 cache_read_input_tokens: 0,
1362 cache_creation_input_tokens: 0,
1363 });
1364 acc.push(&StreamEvent::StopReason("length".into()));
1365
1366 let (text, _tools, usage, stop) = acc.finish_with_usage();
1367 assert_eq!(text, "hello");
1368 let usage = usage.expect("a reported Usage event must not be dropped");
1369 assert_eq!(usage.prompt_tokens, 28);
1370 assert_eq!(usage.completion_tokens, 5);
1371 assert_eq!(usage.total_tokens, 33);
1372 assert_eq!(
1373 stop.as_deref(),
1374 Some("length"),
1375 "the provider stop_reason feeds was_truncated and was being dropped too"
1376 );
1377 }
1378
1379 #[test]
1382 fn absent_usage_stays_none_rather_than_zero() {
1383 let mut acc = StreamAccumulator::default();
1384 acc.push(&StreamEvent::TextDelta("hi".into()));
1385 let (_text, _tools, usage, stop) = acc.finish_with_usage();
1386 assert!(
1387 usage.is_none(),
1388 "no Usage event must yield None so callers can fall back to an estimator"
1389 );
1390 assert!(stop.is_none());
1391 }
1392
1393 #[test]
1401 fn managed_error_events_carry_type_and_code() {
1402 let events = parse_openai_responses_sse_line(
1403 "error",
1404 r#"{"error":{"message":"content refused","type":"invalid_request_error","code":"content_policy_violation"}}"#,
1405 );
1406 let StreamEvent::Error(msg) = events.first().expect("an error event") else {
1407 panic!("expected StreamEvent::Error, got {:?}", events.first());
1408 };
1409 assert!(msg.contains("content refused"), "message dropped: {msg}");
1410 assert!(
1411 msg.contains("type=invalid_request_error"),
1412 "type dropped: {msg}"
1413 );
1414 assert!(
1415 msg.contains("code=content_policy_violation"),
1416 "code dropped: {msg}"
1417 );
1418 }
1419
1420 #[test]
1424 fn a_messageless_managed_error_still_reports_its_code() {
1425 let events = parse_openai_responses_sse_line(
1426 "response.failed",
1427 r#"{"response":{"error":{"code":"content_filter"}}}"#,
1428 );
1429 let StreamEvent::Error(msg) = events.first().expect("an error event") else {
1430 panic!("expected StreamEvent::Error");
1431 };
1432 assert!(msg.contains("managed inference failed"), "{msg}");
1433 assert!(
1434 msg.contains("code=content_filter"),
1435 "classification lost: {msg}"
1436 );
1437 }
1438
1439 #[test]
1442 fn a_bare_managed_error_is_unchanged() {
1443 let events = parse_openai_responses_sse_line("error", r#"{"error":{}}"#);
1444 let StreamEvent::Error(msg) = events.first().expect("an error event") else {
1445 panic!("expected StreamEvent::Error");
1446 };
1447 assert_eq!(msg, "managed inference failed");
1448 }
1449
1450 #[test]
1451 fn parse_openai_text_delta() {
1452 let line = r#"data: {"choices":[{"delta":{"content":"Hello"}}]}"#;
1453 let events = parse_openai_sse_line(line);
1454 assert_eq!(events.len(), 1);
1455 match &events[0] {
1456 StreamEvent::TextDelta(t) => assert_eq!(t, "Hello"),
1457 other => panic!("expected TextDelta, got {:?}", other),
1458 }
1459 }
1460
1461 #[test]
1462 fn parse_openai_tool_call_start() {
1463 let line = r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"edit_file"}}]}}]}"#;
1464 let events = parse_openai_sse_line(line);
1465 assert_eq!(events.len(), 1);
1466 match &events[0] {
1467 StreamEvent::ToolCallStart { name, index, .. } => {
1468 assert_eq!(name, "edit_file");
1469 assert_eq!(*index, 0);
1470 }
1471 other => panic!("expected ToolCallStart, got {:?}", other),
1472 }
1473 }
1474
1475 #[test]
1476 fn parse_openai_tool_call_delta() {
1477 let line = r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"path\":"}}]}}]}"#;
1478 let events = parse_openai_sse_line(line);
1479 assert_eq!(events.len(), 1);
1480 match &events[0] {
1481 StreamEvent::ToolCallDelta {
1482 index,
1483 arguments_delta,
1484 } => {
1485 assert_eq!(*index, 0);
1486 assert!(arguments_delta.contains("path"));
1487 }
1488 other => panic!("expected ToolCallDelta, got {:?}", other),
1489 }
1490 }
1491
1492 #[test]
1493 fn parse_openai_multiple_tool_calls_in_chunk() {
1494 let line = r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"read_file"}},{"index":1,"function":{"name":"search"}}]}}]}"#;
1496 let events = parse_openai_sse_line(line);
1497 assert_eq!(events.len(), 2);
1498 match &events[0] {
1499 StreamEvent::ToolCallStart { name, index, .. } => {
1500 assert_eq!(name, "read_file");
1501 assert_eq!(*index, 0);
1502 }
1503 other => panic!("expected ToolCallStart, got {:?}", other),
1504 }
1505 match &events[1] {
1506 StreamEvent::ToolCallStart { name, index, .. } => {
1507 assert_eq!(name, "search");
1508 assert_eq!(*index, 1);
1509 }
1510 other => panic!("expected ToolCallStart, got {:?}", other),
1511 }
1512 }
1513
1514 #[test]
1515 fn parse_openai_done() {
1516 assert!(parse_openai_sse_line("data: [DONE]").is_empty());
1517 }
1518
1519 #[test]
1520 fn parse_anthropic_text_delta() {
1521 let data = r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"world"}}"#;
1522 let events = parse_anthropic_sse_line("content_block_delta", data);
1523 assert_eq!(events.len(), 1);
1524 match &events[0] {
1525 StreamEvent::TextDelta(t) => assert_eq!(t, "world"),
1526 other => panic!("expected TextDelta, got {:?}", other),
1527 }
1528 }
1529
1530 #[test]
1531 fn parse_anthropic_tool_start() {
1532 let data = r#"{"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"t1","name":"search","input":{}}}"#;
1533 let events = parse_anthropic_sse_line("content_block_start", data);
1534 assert_eq!(events.len(), 1);
1535 match &events[0] {
1536 StreamEvent::ToolCallStart { name, index, .. } => {
1537 assert_eq!(name, "search");
1538 assert_eq!(*index, 1);
1539 }
1540 other => panic!("expected ToolCallStart, got {:?}", other),
1541 }
1542 }
1543
1544 #[test]
1545 fn accumulator_builds_result() {
1546 let mut acc = StreamAccumulator::default();
1547 acc.push(&StreamEvent::TextDelta("Hello ".into()));
1548 acc.push(&StreamEvent::TextDelta("world".into()));
1549 acc.push(&StreamEvent::ToolCallStart {
1550 name: "search".into(),
1551 index: 0,
1552 id: None,
1553 });
1554 acc.push(&StreamEvent::ToolCallDelta {
1555 index: 0,
1556 arguments_delta: r#"{"q":"test"}"#.into(),
1557 });
1558
1559 let (text, tools) = acc.finish();
1560 assert_eq!(text, "Hello world");
1561 assert_eq!(tools.len(), 1);
1562 assert_eq!(tools[0].name, "search");
1563 assert!(tools[0].arguments.contains_key("q"));
1564 }
1565
1566 #[test]
1567 fn parse_sse_lines_openai_format() {
1568 let chunk = "data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\ndata: [DONE]\n\n";
1569 let events = parse_sse_lines(chunk);
1570 assert_eq!(events.len(), 2);
1571 assert_eq!(events[0].0, "message");
1572 assert_eq!(events[1].1, "[DONE]");
1573 }
1574
1575 #[test]
1576 fn parse_sse_lines_anthropic_format() {
1577 let chunk = "event: content_block_delta\ndata: {\"delta\":{\"type\":\"text_delta\",\"text\":\"Hi\"}}\n\n";
1578 let events = parse_sse_lines(chunk);
1579 assert_eq!(events.len(), 1);
1580 assert_eq!(events[0].0, "content_block_delta");
1581 }
1582
1583 #[test]
1584 fn parse_anthropic_message_start_emits_usage() {
1585 let data = r#"{"type":"message_start","message":{"id":"msg_1","role":"assistant","usage":{"input_tokens":245,"output_tokens":1}}}"#;
1586 let events = parse_anthropic_sse_line("message_start", data);
1587 assert_eq!(events.len(), 1);
1588 match &events[0] {
1589 StreamEvent::Usage {
1590 input_tokens,
1591 output_tokens,
1592 ..
1593 } => {
1594 assert_eq!(*input_tokens, 245);
1595 assert_eq!(*output_tokens, 1);
1596 }
1597 other => panic!("expected Usage, got {:?}", other),
1598 }
1599 }
1600
1601 #[test]
1602 fn parse_anthropic_message_delta_emits_stop_reason_and_usage() {
1603 let data = r#"{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":87}}"#;
1604 let events = parse_anthropic_sse_line("message_delta", data);
1605 assert_eq!(events.len(), 2);
1608 match &events[0] {
1609 StreamEvent::StopReason(reason) => assert_eq!(reason, "end_turn"),
1610 other => panic!("expected StopReason, got {:?}", other),
1611 }
1612 match &events[1] {
1613 StreamEvent::Usage {
1614 input_tokens,
1615 output_tokens,
1616 ..
1617 } => {
1618 assert_eq!(*input_tokens, 0);
1619 assert_eq!(*output_tokens, 87);
1620 }
1621 other => panic!("expected Usage, got {:?}", other),
1622 }
1623 }
1624
1625 #[test]
1626 fn parse_anthropic_message_delta_max_tokens_stop_reason() {
1627 let data = r#"{"type":"message_delta","delta":{"stop_reason":"max_tokens"},"usage":{"output_tokens":4096}}"#;
1630 let events = parse_anthropic_sse_line("message_delta", data);
1631 assert!(matches!(
1632 &events[0],
1633 StreamEvent::StopReason(r) if r == "max_tokens"
1634 ));
1635 }
1636
1637 #[test]
1638 fn parse_openai_finish_reason_length_surfaces() {
1639 let line = r#"data: {"choices":[{"delta":{"content":""},"finish_reason":"length"}]}"#;
1642 let events = parse_openai_sse_line(line);
1643 assert!(events
1644 .iter()
1645 .any(|e| matches!(e, StreamEvent::StopReason(r) if r == "length")));
1646 }
1647
1648 #[test]
1649 fn accumulator_captures_stop_reason() {
1650 let mut acc = StreamAccumulator::default();
1651 acc.push(&StreamEvent::TextDelta("partial".into()));
1652 acc.push(&StreamEvent::StopReason("max_tokens".into()));
1653 let (_, _, _, stop) = acc.finish_with_usage();
1654 assert_eq!(stop.as_deref(), Some("max_tokens"));
1655 }
1656
1657 #[test]
1658 fn parse_anthropic_message_start_without_usage_is_empty() {
1659 let data = r#"{"type":"message_start","message":{"id":"msg_1"}}"#;
1661 assert!(parse_anthropic_sse_line("message_start", data).is_empty());
1662 }
1663
1664 #[test]
1665 fn accumulator_tracks_usage_across_anthropic_stream() {
1666 let mut acc = StreamAccumulator::default();
1669 for event in parse_anthropic_sse_line(
1670 "message_start",
1671 r#"{"message":{"usage":{"input_tokens":245,"output_tokens":1}}}"#,
1672 ) {
1673 acc.push(&event);
1674 }
1675 for event in parse_anthropic_sse_line(
1676 "content_block_start",
1677 r#"{"index":0,"content_block":{"type":"text","text":""}}"#,
1678 ) {
1679 acc.push(&event);
1680 }
1681 for (chunk, _) in [
1682 (r#"{"delta":{"type":"text_delta","text":"Hello"}}"#, ()),
1683 (r#"{"delta":{"type":"text_delta","text":", "}}"#, ()),
1684 (r#"{"delta":{"type":"text_delta","text":"world"}}"#, ()),
1685 ] {
1686 for event in parse_anthropic_sse_line("content_block_delta", chunk) {
1687 acc.push(&event);
1688 }
1689 }
1690 for event in parse_anthropic_sse_line("message_delta", r#"{"usage":{"output_tokens":87}}"#)
1691 {
1692 acc.push(&event);
1693 }
1694
1695 let (text, tools, usage, _stop) = acc.finish_with_usage();
1696 assert_eq!(text, "Hello, world");
1697 assert!(tools.is_empty());
1698 let usage = usage.expect("provider reported usage; must surface");
1699 assert_eq!(usage.prompt_tokens, 245);
1700 assert_eq!(usage.completion_tokens, 87);
1702 assert_eq!(usage.total_tokens, 332);
1703 }
1704
1705 #[test]
1706 fn parse_openai_final_chunk_emits_usage() {
1707 let line = r#"data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[],"usage":{"prompt_tokens":245,"completion_tokens":87,"total_tokens":332}}"#;
1710 let events = parse_openai_sse_line(line);
1711 assert_eq!(events.len(), 1);
1712 match &events[0] {
1713 StreamEvent::Usage {
1714 input_tokens,
1715 output_tokens,
1716 ..
1717 } => {
1718 assert_eq!(*input_tokens, 245);
1719 assert_eq!(*output_tokens, 87);
1720 }
1721 other => panic!("expected Usage, got {:?}", other),
1722 }
1723 }
1724
1725 #[test]
1726 fn accumulator_tracks_usage_across_openai_stream() {
1727 let mut acc = StreamAccumulator::default();
1730 for line in [
1731 r#"data: {"choices":[{"delta":{"content":"Hello"}}]}"#,
1732 r#"data: {"choices":[{"delta":{"content":", "}}]}"#,
1733 r#"data: {"choices":[{"delta":{"content":"world"}}]}"#,
1734 r#"data: {"id":"chatcmpl-1","choices":[],"usage":{"prompt_tokens":245,"completion_tokens":87}}"#,
1735 ] {
1736 for event in parse_openai_sse_line(line) {
1737 acc.push(&event);
1738 }
1739 }
1740
1741 let (text, tools, usage, _stop) = acc.finish_with_usage();
1742 assert_eq!(text, "Hello, world");
1743 assert!(tools.is_empty());
1744 let usage = usage.expect("provider reported usage; must surface");
1745 assert_eq!(usage.prompt_tokens, 245);
1746 assert_eq!(usage.completion_tokens, 87);
1747 assert_eq!(usage.total_tokens, 332);
1748 }
1749
1750 #[test]
1751 fn accumulator_returns_no_usage_when_provider_silent() {
1752 let mut acc = StreamAccumulator::default();
1756 acc.push(&StreamEvent::TextDelta("hi".into()));
1757 let (_, _, usage, _stop) = acc.finish_with_usage();
1758 assert!(usage.is_none());
1759 }
1760
1761 #[test]
1762 fn anthropic_stream_decodes_cache_tokens_from_message_start() {
1763 let mut acc = StreamAccumulator::default();
1767 let start = r#"{"message":{"usage":{"input_tokens":50,"output_tokens":1,"cache_read_input_tokens":4000,"cache_creation_input_tokens":600}}}"#;
1768 for e in parse_anthropic_sse_line("message_start", start) {
1769 acc.push(&e);
1770 }
1771 let delta = r#"{"delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":87}}"#;
1772 for e in parse_anthropic_sse_line("message_delta", delta) {
1773 acc.push(&e);
1774 }
1775 let (_t, _c, usage, stop) = acc.finish_with_usage();
1776 let u = usage.expect("usage surfaced");
1777 assert_eq!(u.prompt_tokens, 50, "uncached prefix");
1778 assert_eq!(u.completion_tokens, 87, "final output from message_delta");
1779 assert_eq!(u.cache_read_input_tokens, 4000);
1780 assert_eq!(u.cache_creation_input_tokens, 600);
1781 assert_eq!(stop.as_deref(), Some("end_turn"));
1782 }
1783
1784 #[test]
1785 fn openai_stream_normalizes_cached_tokens_out_of_prompt() {
1786 let mut acc = StreamAccumulator::default();
1790 let chunk = r#"data: {"choices":[],"usage":{"prompt_tokens":1000,"completion_tokens":40,"prompt_tokens_details":{"cached_tokens":800}}}"#;
1791 for e in parse_openai_sse_line(chunk) {
1792 acc.push(&e);
1793 }
1794 let (_t, _c, usage, _s) = acc.finish_with_usage();
1795 let u = usage.expect("usage surfaced");
1796 assert_eq!(u.prompt_tokens, 200, "uncached = 1000 - 800");
1797 assert_eq!(u.cache_read_input_tokens, 800);
1798 assert_eq!(
1799 u.cache_creation_input_tokens, 0,
1800 "OpenAI has no write bucket"
1801 );
1802 }
1803
1804 #[test]
1805 fn responses_failure_is_a_terminal_safe_error_event() {
1806 let events = parse_openai_responses_sse_line(
1807 "response.failed",
1808 r#"{"response":{"error":{"message":"managed model unavailable","stack":"secret"}}}"#,
1809 );
1810 assert!(matches!(
1811 events.as_slice(),
1812 [StreamEvent::Error(message)] if message == "managed model unavailable"
1813 ));
1814 assert!(!format!("{events:?}").contains("secret"));
1815 }
1816
1817 #[test]
1818 fn responses_reasoning_item_done_is_retained_verbatim() {
1819 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"}}"#;
1820 let events = parse_openai_responses_sse_line("response.output_item.done", data);
1821 let expected = serde_json::json!({
1822 "type": "reasoning",
1823 "id": "rs_1",
1824 "status": "completed",
1825 "summary": [{"type": "summary_text", "text": "safe summary"}],
1826 "encrypted_content": "opaque-ciphertext",
1827 });
1828 assert!(matches!(
1829 events.as_slice(),
1830 [StreamEvent::ProviderOutputItem(item)] if item == &expected
1831 ));
1832 let mut accumulator = StreamAccumulator::default();
1833 accumulator.push(&events[0]);
1834 let (_, _, _, _, items) = accumulator.finish_with_provider_output_items();
1835 assert_eq!(items, vec![expected]);
1836 }
1837
1838 #[test]
1839 fn responses_incomplete_is_terminal_failure_not_success() {
1840 let events = parse_openai_responses_sse_line(
1841 "response.incomplete",
1842 r#"{"response":{"status":"incomplete","incomplete_details":{"reason":"max_output_tokens"},"usage":{"input_tokens":17,"output_tokens":9}}}"#,
1843 );
1844 assert!(
1845 events
1846 .iter()
1847 .any(|event| matches!(event, StreamEvent::Error(_))),
1848 "response.incomplete must emit a terminal error"
1849 );
1850 }
1851
1852 fn run_stop_filter(stops: &[&str], events: Vec<StreamEvent>) -> Vec<StreamEvent> {
1853 let mut filter = StopSequenceFilter::new(stops.iter().map(|s| s.to_string()).collect());
1854 let mut out = Vec::new();
1855 for event in events {
1856 out.extend(filter.push(event));
1857 }
1858 out.extend(filter.finish());
1859 out
1860 }
1861
1862 fn streamed_text(events: &[StreamEvent]) -> String {
1863 events
1864 .iter()
1865 .filter_map(|event| match event {
1866 StreamEvent::TextDelta(delta) => Some(delta.as_str()),
1867 _ => None,
1868 })
1869 .collect()
1870 }
1871
1872 #[test]
1873 fn stop_filter_catches_a_sequence_split_across_three_deltas() {
1874 let out = run_stop_filter(
1875 &["STOPHERE"],
1876 vec![
1877 StreamEvent::TextDelta("Paris S".into()),
1878 StreamEvent::TextDelta("TOPH".into()),
1879 StreamEvent::Usage {
1880 input_tokens: 3,
1881 output_tokens: 2,
1882 cache_read_input_tokens: 0,
1883 cache_creation_input_tokens: 0,
1884 },
1885 StreamEvent::TextDelta("ERE indeed".into()),
1886 StreamEvent::TextDelta(" and more".into()),
1887 StreamEvent::StopReason("length".into()),
1888 StreamEvent::Done {
1889 text: "Paris STOPHERE indeed and more".into(),
1890 tool_calls: Vec::new(),
1891 },
1892 ],
1893 );
1894 assert_eq!(streamed_text(&out), "Paris ");
1895 let reasons: Vec<&str> = out
1896 .iter()
1897 .filter_map(|event| match event {
1898 StreamEvent::StopReason(reason) => Some(reason.as_str()),
1899 _ => None,
1900 })
1901 .collect();
1902 assert_eq!(reasons, ["stop_sequence"]);
1903 assert!(matches!(
1904 out.last(),
1905 Some(StreamEvent::Done { text, .. }) if text == "Paris "
1906 ));
1907 assert!(out
1909 .iter()
1910 .any(|event| matches!(event, StreamEvent::Usage { .. })));
1911 }
1912
1913 #[test]
1916 fn stop_filter_releases_a_false_alarm_tail_unchanged() {
1917 let out = run_stop_filter(
1918 &["STOPHERE"],
1919 vec![
1920 StreamEvent::TextDelta("almost STO".into()),
1921 StreamEvent::TextDelta("P but not quite".into()),
1922 StreamEvent::StopReason("stop".into()),
1923 StreamEvent::Done {
1924 text: "almost STOP but not quite".into(),
1925 tool_calls: Vec::new(),
1926 },
1927 ],
1928 );
1929 assert_eq!(streamed_text(&out), "almost STOP but not quite");
1930 assert!(out
1931 .iter()
1932 .any(|event| matches!(event, StreamEvent::StopReason(reason) if reason == "stop")));
1933 assert!(matches!(
1934 out.last(),
1935 Some(StreamEvent::Done { text, .. }) if text == "almost STOP but not quite"
1936 ));
1937 }
1938
1939 #[test]
1943 fn stop_filter_reports_one_reason_when_only_done_carries_the_stop() {
1944 let out = run_stop_filter(
1945 &["STOPHERE"],
1946 vec![
1947 StreamEvent::StopReason("length".into()),
1948 StreamEvent::Done {
1949 text: "Paris STOPHERE indeed".into(),
1950 tool_calls: Vec::new(),
1951 },
1952 ],
1953 );
1954 let reasons: Vec<&str> = out
1955 .iter()
1956 .filter_map(|event| match event {
1957 StreamEvent::StopReason(reason) => Some(reason.as_str()),
1958 _ => None,
1959 })
1960 .collect();
1961 assert_eq!(reasons, ["stop_sequence"]);
1962 assert!(matches!(
1963 out.last(),
1964 Some(StreamEvent::Done { text, .. }) if text == "Paris "
1965 ));
1966 }
1967
1968 #[test]
1969 fn stop_filter_flushes_a_held_tail_when_the_source_closes() {
1970 let out = run_stop_filter(
1971 &["STOPHERE"],
1972 vec![StreamEvent::TextDelta("ends ST".into())],
1973 );
1974 assert_eq!(streamed_text(&out), "ends ST");
1975 }
1976
1977 #[test]
1978 fn stop_filter_holds_back_only_at_char_boundaries() {
1979 let out = run_stop_filter(
1980 &["\u{e9}\u{e9}"],
1981 vec![
1982 StreamEvent::TextDelta("caf\u{e9}".into()),
1983 StreamEvent::TextDelta("\u{e9} au lait".into()),
1984 ],
1985 );
1986 assert_eq!(streamed_text(&out), "caf");
1987 }
1988
1989 #[test]
1990 fn apply_stop_backstop_cuts_text_and_reports_stop_sequence() {
1991 use crate::tasks::generate::apply_stop_backstop;
1992 let stops = vec!["STOPHERE".to_string()];
1993 assert_eq!(
1994 apply_stop_backstop(
1995 "Paris STOPHERE indeed".into(),
1996 Some("length".into()),
1997 &stops,
1998 0
1999 ),
2000 ("Paris ".to_string(), Some("stop_sequence".to_string()))
2001 );
2002 assert_eq!(
2004 apply_stop_backstop("Paris".into(), Some("stop".into()), &stops, 0),
2005 ("Paris".to_string(), Some("stop".to_string()))
2006 );
2007 assert_eq!(
2009 apply_stop_backstop("a STOPHERE b".into(), None, &[], 0),
2010 ("a STOPHERE b".to_string(), None)
2011 );
2012 }
2013}