1use agent_sdk_foundation::llm::{ContentBlock, StopReason, Usage};
8#[cfg(any(feature = "openai", feature = "openai-codex"))]
9use bytes::BytesMut;
10use futures::Stream;
11use std::collections::HashMap;
12use std::pin::Pin;
13
14const MAX_BLOCK_INDEX: usize = 4096;
23
24#[cfg(any(feature = "openai", feature = "openai-codex"))]
40#[derive(Debug, Default)]
41pub(crate) struct SseLineBuffer {
42 buf: BytesMut,
43}
44
45#[cfg(any(feature = "openai", feature = "openai-codex"))]
46impl SseLineBuffer {
47 #[must_use]
49 pub(crate) fn new() -> Self {
50 Self::default()
51 }
52
53 pub(crate) fn extend(&mut self, chunk: &[u8]) {
55 self.buf.extend_from_slice(chunk);
56 }
57
58 pub(crate) fn next_line(&mut self) -> Option<String> {
63 let newline = self.buf.iter().position(|&b| b == b'\n')?;
64 let mut line = self.buf.split_to(newline + 1);
65 line.truncate(newline);
66 Some(String::from_utf8_lossy(&line).into_owned())
67 }
68}
69
70#[derive(Debug, Clone)]
75#[non_exhaustive]
76pub enum StreamDelta {
77 TextDelta {
79 delta: String,
81 block_index: usize,
83 },
84
85 ThinkingDelta {
87 delta: String,
89 block_index: usize,
91 },
92
93 ToolUseStart {
95 id: String,
97 name: String,
99 block_index: usize,
101 thought_signature: Option<String>,
103 },
104
105 ToolInputDelta {
107 id: String,
109 delta: String,
111 block_index: usize,
113 },
114
115 Usage(Usage),
117
118 Done {
120 stop_reason: Option<StopReason>,
122 },
123
124 SignatureDelta {
126 delta: String,
128 block_index: usize,
130 },
131
132 RedactedThinking {
134 data: String,
136 block_index: usize,
138 },
139
140 OpaqueReasoning {
146 provider: String,
148 data: serde_json::Value,
150 block_index: usize,
152 },
153
154 Error {
156 message: String,
158 kind: StreamErrorKind,
163 },
164}
165
166#[derive(Debug, Clone, Copy, PartialEq, Eq)]
175#[non_exhaustive]
176pub enum StreamErrorKind {
177 RateLimited,
179 ServerError,
182 InvalidRequest,
185 Unknown,
195}
196
197impl StreamErrorKind {
198 #[must_use]
202 pub const fn is_recoverable(self) -> bool {
203 matches!(self, Self::RateLimited | Self::ServerError)
204 }
205}
206
207pub type StreamBox<'a> = Pin<Box<dyn Stream<Item = anyhow::Result<StreamDelta>> + Send + 'a>>;
209
210#[derive(Debug, Default)]
215pub struct StreamAccumulator {
216 text_blocks: Vec<String>,
218 thinking_blocks: Vec<String>,
220 thinking_signatures: HashMap<usize, String>,
222 redacted_thinking_blocks: Vec<(usize, String)>,
224 opaque_reasoning_blocks: Vec<(usize, String, serde_json::Value)>,
226 tool_uses: Vec<ToolUseAccumulator>,
228 usage: Option<Usage>,
230 stop_reason: Option<StopReason>,
232}
233
234#[derive(Debug, Default)]
236pub struct ToolUseAccumulator {
237 pub id: String,
239 pub name: String,
241 pub input_json: String,
243 pub block_index: usize,
245 pub thought_signature: Option<String>,
247}
248
249impl StreamAccumulator {
250 #[must_use]
252 pub fn new() -> Self {
253 Self::default()
254 }
255
256 pub fn apply(&mut self, delta: &StreamDelta) {
258 match delta {
259 StreamDelta::TextDelta { delta, block_index } => {
260 if *block_index > MAX_BLOCK_INDEX {
261 log::warn!(
262 "dropping TextDelta with out-of-range block_index {block_index} (max {MAX_BLOCK_INDEX})"
263 );
264 return;
265 }
266 while self.text_blocks.len() <= *block_index {
267 self.text_blocks.push(String::new());
268 }
269 self.text_blocks[*block_index].push_str(delta);
270 }
271 StreamDelta::ThinkingDelta { delta, block_index } => {
272 if *block_index > MAX_BLOCK_INDEX {
273 log::warn!(
274 "dropping ThinkingDelta with out-of-range block_index {block_index} (max {MAX_BLOCK_INDEX})"
275 );
276 return;
277 }
278 while self.thinking_blocks.len() <= *block_index {
279 self.thinking_blocks.push(String::new());
280 }
281 self.thinking_blocks[*block_index].push_str(delta);
282 }
283 StreamDelta::ToolUseStart {
284 id,
285 name,
286 block_index,
287 thought_signature,
288 } => {
289 self.tool_uses.push(ToolUseAccumulator {
290 id: id.clone(),
291 name: name.clone(),
292 input_json: String::new(),
293 block_index: *block_index,
294 thought_signature: thought_signature.clone(),
295 });
296 }
297 StreamDelta::ToolInputDelta { id, delta, .. } => {
298 if let Some(tool) = self.tool_uses.iter_mut().find(|t| t.id == *id) {
299 tool.input_json.push_str(delta);
300 }
301 }
302 StreamDelta::SignatureDelta { delta, block_index } => {
303 self.thinking_signatures
304 .entry(*block_index)
305 .or_default()
306 .push_str(delta);
307 }
308 StreamDelta::RedactedThinking { data, block_index } => {
309 self.redacted_thinking_blocks
310 .push((*block_index, data.clone()));
311 }
312 StreamDelta::OpaqueReasoning {
313 provider,
314 data,
315 block_index,
316 } => {
317 self.opaque_reasoning_blocks
318 .push((*block_index, provider.clone(), data.clone()));
319 }
320 StreamDelta::Usage(u) => {
321 self.usage = Some(u.clone());
322 }
323 StreamDelta::Done { stop_reason } => {
324 self.stop_reason = *stop_reason;
325 }
326 StreamDelta::Error { .. } => {}
327 }
328 }
329
330 #[must_use]
332 pub const fn usage(&self) -> Option<&Usage> {
333 self.usage.as_ref()
334 }
335
336 #[must_use]
338 pub const fn stop_reason(&self) -> Option<&StopReason> {
339 self.stop_reason.as_ref()
340 }
341
342 #[must_use]
347 pub fn into_content_blocks(self) -> Vec<ContentBlock> {
348 let mut blocks: Vec<(usize, ContentBlock)> = Vec::new();
349
350 let mut signatures = self.thinking_signatures;
352 for (idx, thinking) in self.thinking_blocks.into_iter().enumerate() {
353 if !thinking.is_empty() {
354 let signature = signatures.remove(&idx).filter(|s| !s.is_empty());
355 blocks.push((
356 idx,
357 ContentBlock::Thinking {
358 thinking,
359 signature,
360 },
361 ));
362 }
363 }
364
365 for (idx, data) in self.redacted_thinking_blocks {
367 blocks.push((idx, ContentBlock::RedactedThinking { data }));
368 }
369
370 for (idx, provider, data) in self.opaque_reasoning_blocks {
372 blocks.push((idx, ContentBlock::OpaqueReasoning { provider, data }));
373 }
374
375 for (idx, text) in self.text_blocks.into_iter().enumerate() {
377 if !text.is_empty() {
378 blocks.push((idx, ContentBlock::Text { text }));
379 }
380 }
381
382 for tool in self.tool_uses {
384 let input: serde_json::Value =
385 serde_json::from_str(&tool.input_json).unwrap_or_else(|e| {
386 log::warn!(
387 "Failed to parse streamed tool input JSON for tool '{}' (id={}): {} — \
388 input_json ({} bytes): '{}'",
389 tool.name,
390 tool.id,
391 e,
392 tool.input_json.len(),
393 tool.input_json.chars().take(500).collect::<String>(),
394 );
395 serde_json::json!({})
396 });
397 blocks.push((
398 tool.block_index,
399 ContentBlock::ToolUse {
400 id: tool.id,
401 name: tool.name,
402 input,
403 thought_signature: tool.thought_signature,
404 },
405 ));
406 }
407
408 blocks.sort_by_key(|(idx, _)| *idx);
410
411 blocks.into_iter().map(|(_, block)| block).collect()
412 }
413
414 pub const fn take_usage(&mut self) -> Option<Usage> {
416 self.usage.take()
417 }
418
419 pub const fn take_stop_reason(&mut self) -> Option<StopReason> {
421 self.stop_reason.take()
422 }
423}
424
425#[cfg(test)]
426mod tests {
427 use super::*;
428
429 #[test]
430 fn test_accumulator_text_deltas() {
431 let mut acc = StreamAccumulator::new();
432
433 acc.apply(&StreamDelta::TextDelta {
434 delta: "Hello".to_string(),
435 block_index: 0,
436 });
437 acc.apply(&StreamDelta::TextDelta {
438 delta: " world".to_string(),
439 block_index: 0,
440 });
441
442 let blocks = acc.into_content_blocks();
443 assert_eq!(blocks.len(), 1);
444 assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "Hello world"));
445 }
446
447 #[test]
448 fn test_accumulator_multiple_text_blocks() {
449 let mut acc = StreamAccumulator::new();
450
451 acc.apply(&StreamDelta::TextDelta {
452 delta: "First".to_string(),
453 block_index: 0,
454 });
455 acc.apply(&StreamDelta::TextDelta {
456 delta: "Second".to_string(),
457 block_index: 1,
458 });
459
460 let blocks = acc.into_content_blocks();
461 assert_eq!(blocks.len(), 2);
462 assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "First"));
463 assert!(matches!(&blocks[1], ContentBlock::Text { text } if text == "Second"));
464 }
465
466 #[test]
467 fn test_accumulator_thinking_signature() {
468 let mut acc = StreamAccumulator::new();
469
470 acc.apply(&StreamDelta::ThinkingDelta {
471 delta: "Reasoning".to_string(),
472 block_index: 0,
473 });
474 acc.apply(&StreamDelta::SignatureDelta {
475 delta: "sig_123".to_string(),
476 block_index: 0,
477 });
478
479 let blocks = acc.into_content_blocks();
480 assert_eq!(blocks.len(), 1);
481 assert!(matches!(
482 &blocks[0],
483 ContentBlock::Thinking { thinking, signature }
484 if thinking == "Reasoning" && signature.as_deref() == Some("sig_123")
485 ));
486 }
487
488 #[test]
489 fn accumulator_preserves_opaque_reasoning_payload_and_order() {
490 let mut acc = StreamAccumulator::new();
491 acc.apply(&StreamDelta::TextDelta {
492 delta: "visible".to_owned(),
493 block_index: 2,
494 });
495 acc.apply(&StreamDelta::OpaqueReasoning {
496 provider: "test-provider".to_owned(),
497 data: serde_json::json!({
498 "id": "reasoning_1",
499 "encrypted_content": "do-not-inspect"
500 }),
501 block_index: 1,
502 });
503
504 let blocks = acc.into_content_blocks();
505 assert_eq!(blocks.len(), 2);
506 assert!(matches!(
507 &blocks[0],
508 ContentBlock::OpaqueReasoning { provider, data }
509 if provider == "test-provider"
510 && data["id"] == "reasoning_1"
511 && data["encrypted_content"] == "do-not-inspect"
512 ));
513 assert!(matches!(
514 &blocks[1],
515 ContentBlock::Text { text } if text == "visible"
516 ));
517 }
518
519 #[test]
520 fn test_accumulator_tool_use() {
521 let mut acc = StreamAccumulator::new();
522
523 acc.apply(&StreamDelta::ToolUseStart {
524 id: "call_123".to_string(),
525 name: "read_file".to_string(),
526 block_index: 0,
527 thought_signature: None,
528 });
529 acc.apply(&StreamDelta::ToolInputDelta {
530 id: "call_123".to_string(),
531 delta: r#"{"path":"#.to_string(),
532 block_index: 0,
533 });
534 acc.apply(&StreamDelta::ToolInputDelta {
535 id: "call_123".to_string(),
536 delta: r#""test.txt"}"#.to_string(),
537 block_index: 0,
538 });
539
540 let blocks = acc.into_content_blocks();
541 assert_eq!(blocks.len(), 1);
542 match &blocks[0] {
543 ContentBlock::ToolUse {
544 id, name, input, ..
545 } => {
546 assert_eq!(id, "call_123");
547 assert_eq!(name, "read_file");
548 assert_eq!(input["path"], "test.txt");
549 }
550 _ => panic!("Expected ToolUse block"),
551 }
552 }
553
554 #[test]
555 fn test_accumulator_mixed_content() {
556 let mut acc = StreamAccumulator::new();
557
558 acc.apply(&StreamDelta::TextDelta {
559 delta: "Let me read that file.".to_string(),
560 block_index: 0,
561 });
562 acc.apply(&StreamDelta::ToolUseStart {
563 id: "call_456".to_string(),
564 name: "read_file".to_string(),
565 block_index: 1,
566 thought_signature: None,
567 });
568 acc.apply(&StreamDelta::ToolInputDelta {
569 id: "call_456".to_string(),
570 delta: r#"{"path":"file.txt"}"#.to_string(),
571 block_index: 1,
572 });
573 acc.apply(&StreamDelta::Usage(Usage {
574 input_tokens: 100,
575 output_tokens: 50,
576 cached_input_tokens: 0,
577 cache_creation_input_tokens: 0,
578 }));
579 acc.apply(&StreamDelta::Done {
580 stop_reason: Some(StopReason::ToolUse),
581 });
582
583 assert!(acc.usage().is_some());
584 assert_eq!(acc.usage().map(|u| u.input_tokens), Some(100));
585 assert!(matches!(acc.stop_reason(), Some(StopReason::ToolUse)));
586
587 let blocks = acc.into_content_blocks();
588 assert_eq!(blocks.len(), 2);
589 assert!(matches!(&blocks[0], ContentBlock::Text { .. }));
590 assert!(matches!(&blocks[1], ContentBlock::ToolUse { .. }));
591 }
592
593 #[test]
594 fn test_accumulator_invalid_tool_json() {
595 let mut acc = StreamAccumulator::new();
596
597 acc.apply(&StreamDelta::ToolUseStart {
598 id: "call_789".to_string(),
599 name: "test_tool".to_string(),
600 block_index: 0,
601 thought_signature: None,
602 });
603 acc.apply(&StreamDelta::ToolInputDelta {
604 id: "call_789".to_string(),
605 delta: "invalid json {".to_string(),
606 block_index: 0,
607 });
608
609 let blocks = acc.into_content_blocks();
610 assert_eq!(blocks.len(), 1);
611 match &blocks[0] {
612 ContentBlock::ToolUse { input, .. } => {
613 assert!(input.is_object());
614 }
615 _ => panic!("Expected ToolUse block"),
616 }
617 }
618
619 #[test]
620 fn test_accumulator_empty_tool_input_falls_back_to_empty_object() {
621 let mut acc = StreamAccumulator::new();
626
627 acc.apply(&StreamDelta::ToolUseStart {
628 id: "call_empty".to_string(),
629 name: "read".to_string(),
630 block_index: 0,
631 thought_signature: None,
632 });
633 let blocks = acc.into_content_blocks();
636 assert_eq!(blocks.len(), 1);
637 match &blocks[0] {
638 ContentBlock::ToolUse { input, name, .. } => {
639 assert_eq!(name, "read");
640 assert_eq!(input, &serde_json::json!({}));
641 }
642 _ => panic!("Expected ToolUse block"),
643 }
644 }
645
646 #[test]
647 fn test_accumulator_mismatched_delta_id_drops_input() {
648 let mut acc = StreamAccumulator::new();
651
652 acc.apply(&StreamDelta::ToolUseStart {
653 id: "call_A".to_string(),
654 name: "bash".to_string(),
655 block_index: 0,
656 thought_signature: None,
657 });
658 acc.apply(&StreamDelta::ToolInputDelta {
660 id: "call_B".to_string(),
661 delta: r#"{"command":"ls"}"#.to_string(),
662 block_index: 0,
663 });
664
665 let blocks = acc.into_content_blocks();
666 assert_eq!(blocks.len(), 1);
667 match &blocks[0] {
668 ContentBlock::ToolUse { input, .. } => {
669 assert_eq!(input, &serde_json::json!({}));
671 }
672 _ => panic!("Expected ToolUse block"),
673 }
674 }
675
676 #[test]
677 fn test_accumulator_empty() {
678 let acc = StreamAccumulator::new();
679 let blocks = acc.into_content_blocks();
680 assert!(blocks.is_empty());
681 }
682
683 #[test]
684 fn test_accumulator_skips_empty_text() {
685 let mut acc = StreamAccumulator::new();
686
687 acc.apply(&StreamDelta::TextDelta {
688 delta: String::new(),
689 block_index: 0,
690 });
691 acc.apply(&StreamDelta::TextDelta {
692 delta: "Hello".to_string(),
693 block_index: 1,
694 });
695
696 let blocks = acc.into_content_blocks();
697 assert_eq!(blocks.len(), 1);
698 assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "Hello"));
699 }
700
701 #[test]
702 fn test_accumulator_ignores_out_of_range_block_index() {
703 let mut acc = StreamAccumulator::new();
707
708 acc.apply(&StreamDelta::TextDelta {
709 delta: "ok".to_string(),
710 block_index: 0,
711 });
712 acc.apply(&StreamDelta::TextDelta {
713 delta: "boom".to_string(),
714 block_index: usize::MAX,
715 });
716 acc.apply(&StreamDelta::ThinkingDelta {
717 delta: "boom".to_string(),
718 block_index: usize::MAX,
719 });
720
721 let blocks = acc.into_content_blocks();
722 assert_eq!(blocks.len(), 1);
723 assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "ok"));
724 }
725
726 #[cfg(any(feature = "openai", feature = "openai-codex"))]
727 #[test]
728 fn test_sse_line_buffer_splits_multiple_lines() {
729 let mut buf = SseLineBuffer::new();
730 buf.extend(b"data: one\ndata: two\n");
731 assert_eq!(buf.next_line().as_deref(), Some("data: one"));
732 assert_eq!(buf.next_line().as_deref(), Some("data: two"));
733 assert_eq!(buf.next_line(), None);
734 }
735
736 #[cfg(any(feature = "openai", feature = "openai-codex"))]
737 #[test]
738 fn test_sse_line_buffer_buffers_partial_line_until_newline() {
739 let mut buf = SseLineBuffer::new();
740 buf.extend(b"data: par");
741 assert_eq!(buf.next_line(), None);
742 buf.extend(b"tial\n");
743 assert_eq!(buf.next_line().as_deref(), Some("data: partial"));
744 }
745
746 #[cfg(any(feature = "openai", feature = "openai-codex"))]
747 #[test]
748 fn test_sse_line_buffer_handles_utf8_split_across_chunks() {
749 let mut buf = SseLineBuffer::new();
754 let line = "data: café\n";
755 let bytes = line.as_bytes();
756 let split = bytes.len() - 2; buf.extend(&bytes[..split]);
758 assert_eq!(buf.next_line(), None);
759 buf.extend(&bytes[split..]);
760 assert_eq!(buf.next_line().as_deref(), Some("data: café"));
761 }
762}