agent_sdk/llm/

streaming.rs

//! Streaming types for LLM responses.
//!
//! This module provides types for handling streaming responses from LLM providers.
//! The [`StreamDelta`] enum represents individual events in a streaming response,
//! and [`StreamAccumulator`] helps collect these events into a final response.

use crate::llm::{ContentBlock, StopReason, Usage};
use futures::Stream;
use std::collections::HashMap;
use std::pin::Pin;

/// Events yielded during streaming LLM responses.
///
/// Each variant represents a different type of event that can occur
/// during a streaming response from an LLM provider.
#[derive(Debug, Clone)]
pub enum StreamDelta {
    /// A text delta for streaming text content.
    TextDelta {
        /// The text fragment to append
        delta: String,
        /// Index of the content block being streamed
        block_index: usize,
    },

    /// A thinking delta for streaming thinking/reasoning content.
    ThinkingDelta {
        /// The thinking fragment to append
        delta: String,
        /// Index of the content block being streamed
        block_index: usize,
    },

    /// Start of a tool use block (name and id are known).
    ToolUseStart {
        /// Unique identifier for this tool call
        id: String,
        /// Name of the tool being called
        name: String,
        /// Index of the content block
        block_index: usize,
        /// Optional thought signature (used by Gemini 3.x models)
        thought_signature: Option<String>,
    },

    /// Incremental JSON for tool input (partial/incomplete JSON).
    ToolInputDelta {
        /// Tool call ID this delta belongs to
        id: String,
        /// JSON fragment to append
        delta: String,
        /// Index of the content block
        block_index: usize,
    },

    /// Usage information (typically at stream end).
    Usage(Usage),

    /// Stream completed with stop reason.
    Done {
        /// Why the stream ended
        stop_reason: Option<StopReason>,
    },

    /// A signature delta for a thinking block.
    SignatureDelta {
        /// The signature fragment to append
        delta: String,
        /// Index of the content block being streamed
        block_index: usize,
    },

    /// A redacted thinking block received at `content_block_start`.
    RedactedThinking {
        /// Opaque data payload
        data: String,
        /// Index of the content block
        block_index: usize,
    },

    /// Error during streaming.
    Error {
        /// Error message
        message: String,
        /// Whether the error is recoverable (e.g., rate limit)
        recoverable: bool,
    },
}

/// Type alias for a boxed stream of stream deltas.
pub type StreamBox<'a> = Pin<Box<dyn Stream<Item = anyhow::Result<StreamDelta>> + Send + 'a>>;

/// Helper to accumulate streamed content into a final response.
///
/// This struct collects [`StreamDelta`] events and can convert them
/// into the final content blocks once the stream is complete.
#[derive(Debug, Default)]
pub struct StreamAccumulator {
    /// Accumulated text for each block index
    text_blocks: Vec<String>,
    /// Accumulated thinking blocks for each block index
    thinking_blocks: Vec<String>,
    /// Accumulated signatures keyed by block index
    thinking_signatures: HashMap<usize, String>,
    /// Redacted thinking blocks: (`block_index`, data)
    redacted_thinking_blocks: Vec<(usize, String)>,
    /// Accumulated tool use calls
    tool_uses: Vec<ToolUseAccumulator>,
    /// Usage information from the stream
    usage: Option<Usage>,
    /// Stop reason from the stream
    stop_reason: Option<StopReason>,
}

/// Accumulator for a single tool use during streaming.
#[derive(Debug, Default)]
pub struct ToolUseAccumulator {
    /// Tool call ID
    pub id: String,
    /// Tool name
    pub name: String,
    /// Accumulated JSON input (may be incomplete during streaming)
    pub input_json: String,
    /// Block index for ordering
    pub block_index: usize,
    /// Optional thought signature (used by Gemini 3.x models)
    pub thought_signature: Option<String>,
}

impl StreamAccumulator {
    /// Create a new empty accumulator.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Apply a stream delta to the accumulator.
    pub fn apply(&mut self, delta: &StreamDelta) {
        match delta {
            StreamDelta::TextDelta { delta, block_index } => {
                while self.text_blocks.len() <= *block_index {
                    self.text_blocks.push(String::new());
                }
                self.text_blocks[*block_index].push_str(delta);
            }
            StreamDelta::ThinkingDelta { delta, block_index } => {
                while self.thinking_blocks.len() <= *block_index {
                    self.thinking_blocks.push(String::new());
                }
                self.thinking_blocks[*block_index].push_str(delta);
            }
            StreamDelta::ToolUseStart {
                id,
                name,
                block_index,
                thought_signature,
            } => {
                self.tool_uses.push(ToolUseAccumulator {
                    id: id.clone(),
                    name: name.clone(),
                    input_json: String::new(),
                    block_index: *block_index,
                    thought_signature: thought_signature.clone(),
                });
            }
            StreamDelta::ToolInputDelta { id, delta, .. } => {
                if let Some(tool) = self.tool_uses.iter_mut().find(|t| t.id == *id) {
                    tool.input_json.push_str(delta);
                }
            }
            StreamDelta::SignatureDelta { delta, block_index } => {
                self.thinking_signatures
                    .entry(*block_index)
                    .or_default()
                    .push_str(delta);
            }
            StreamDelta::RedactedThinking { data, block_index } => {
                self.redacted_thinking_blocks
                    .push((*block_index, data.clone()));
            }
            StreamDelta::Usage(u) => {
                self.usage = Some(u.clone());
            }
            StreamDelta::Done { stop_reason } => {
                self.stop_reason = *stop_reason;
            }
            StreamDelta::Error { .. } => {}
        }
    }

    /// Get the accumulated usage information.
    #[must_use]
    pub const fn usage(&self) -> Option<&Usage> {
        self.usage.as_ref()
    }

    /// Get the stop reason.
    #[must_use]
    pub const fn stop_reason(&self) -> Option<&StopReason> {
        self.stop_reason.as_ref()
    }

    /// Convert accumulated content to `ContentBlock`s.
    ///
    /// This consumes the accumulator and returns the final content blocks.
    /// Tool use JSON is parsed at this point; invalid JSON results in a null input.
    #[must_use]
    pub fn into_content_blocks(self) -> Vec<ContentBlock> {
        let mut blocks: Vec<(usize, ContentBlock)> = Vec::new();

        // Add thinking blocks with their indices, attaching signatures
        let mut signatures = self.thinking_signatures;
        for (idx, thinking) in self.thinking_blocks.into_iter().enumerate() {
            if !thinking.is_empty() {
                let signature = signatures.remove(&idx).filter(|s| !s.is_empty());
                blocks.push((
                    idx,
                    ContentBlock::Thinking {
                        thinking,
                        signature,
                    },
                ));
            }
        }

        // Add redacted thinking blocks
        for (idx, data) in self.redacted_thinking_blocks {
            blocks.push((idx, ContentBlock::RedactedThinking { data }));
        }

        // Add text blocks with their indices
        for (idx, text) in self.text_blocks.into_iter().enumerate() {
            if !text.is_empty() {
                blocks.push((idx, ContentBlock::Text { text }));
            }
        }

        // Add tool uses with their indices
        for tool in self.tool_uses {
            let input: serde_json::Value =
                serde_json::from_str(&tool.input_json).unwrap_or_else(|e| {
                    log::warn!(
                        "Failed to parse streamed tool input JSON for tool '{}' (id={}): {} — \
                         input_json ({} bytes): '{}'",
                        tool.name,
                        tool.id,
                        e,
                        tool.input_json.len(),
                        tool.input_json.chars().take(500).collect::<String>(),
                    );
                    serde_json::json!({})
                });
            blocks.push((
                tool.block_index,
                ContentBlock::ToolUse {
                    id: tool.id,
                    name: tool.name,
                    input,
                    thought_signature: tool.thought_signature,
                },
            ));
        }

        // Sort by block index to maintain order
        blocks.sort_by_key(|(idx, _)| *idx);

        blocks.into_iter().map(|(_, block)| block).collect()
    }

    /// Take ownership of accumulated usage.
    pub const fn take_usage(&mut self) -> Option<Usage> {
        self.usage.take()
    }

    /// Take ownership of stop reason.
    pub const fn take_stop_reason(&mut self) -> Option<StopReason> {
        self.stop_reason.take()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_accumulator_text_deltas() {
        let mut acc = StreamAccumulator::new();

        acc.apply(&StreamDelta::TextDelta {
            delta: "Hello".to_string(),
            block_index: 0,
        });
        acc.apply(&StreamDelta::TextDelta {
            delta: " world".to_string(),
            block_index: 0,
        });

        let blocks = acc.into_content_blocks();
        assert_eq!(blocks.len(), 1);
        assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "Hello world"));
    }

    #[test]
    fn test_accumulator_multiple_text_blocks() {
        let mut acc = StreamAccumulator::new();

        acc.apply(&StreamDelta::TextDelta {
            delta: "First".to_string(),
            block_index: 0,
        });
        acc.apply(&StreamDelta::TextDelta {
            delta: "Second".to_string(),
            block_index: 1,
        });

        let blocks = acc.into_content_blocks();
        assert_eq!(blocks.len(), 2);
        assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "First"));
        assert!(matches!(&blocks[1], ContentBlock::Text { text } if text == "Second"));
    }

    #[test]
    fn test_accumulator_thinking_signature() {
        let mut acc = StreamAccumulator::new();

        acc.apply(&StreamDelta::ThinkingDelta {
            delta: "Reasoning".to_string(),
            block_index: 0,
        });
        acc.apply(&StreamDelta::SignatureDelta {
            delta: "sig_123".to_string(),
            block_index: 0,
        });

        let blocks = acc.into_content_blocks();
        assert_eq!(blocks.len(), 1);
        assert!(matches!(
            &blocks[0],
            ContentBlock::Thinking { thinking, signature }
            if thinking == "Reasoning" && signature.as_deref() == Some("sig_123")
        ));
    }

    #[test]
    fn test_accumulator_tool_use() {
        let mut acc = StreamAccumulator::new();

        acc.apply(&StreamDelta::ToolUseStart {
            id: "call_123".to_string(),
            name: "read_file".to_string(),
            block_index: 0,
            thought_signature: None,
        });
        acc.apply(&StreamDelta::ToolInputDelta {
            id: "call_123".to_string(),
            delta: r#"{"path":"#.to_string(),
            block_index: 0,
        });
        acc.apply(&StreamDelta::ToolInputDelta {
            id: "call_123".to_string(),
            delta: r#""test.txt"}"#.to_string(),
            block_index: 0,
        });

        let blocks = acc.into_content_blocks();
        assert_eq!(blocks.len(), 1);
        match &blocks[0] {
            ContentBlock::ToolUse {
                id, name, input, ..
            } => {
                assert_eq!(id, "call_123");
                assert_eq!(name, "read_file");
                assert_eq!(input["path"], "test.txt");
            }
            _ => panic!("Expected ToolUse block"),
        }
    }

    #[test]
    fn test_accumulator_mixed_content() {
        let mut acc = StreamAccumulator::new();

        acc.apply(&StreamDelta::TextDelta {
            delta: "Let me read that file.".to_string(),
            block_index: 0,
        });
        acc.apply(&StreamDelta::ToolUseStart {
            id: "call_456".to_string(),
            name: "read_file".to_string(),
            block_index: 1,
            thought_signature: None,
        });
        acc.apply(&StreamDelta::ToolInputDelta {
            id: "call_456".to_string(),
            delta: r#"{"path":"file.txt"}"#.to_string(),
            block_index: 1,
        });
        acc.apply(&StreamDelta::Usage(Usage {
            input_tokens: 100,
            output_tokens: 50,
            cached_input_tokens: 0,
        }));
        acc.apply(&StreamDelta::Done {
            stop_reason: Some(StopReason::ToolUse),
        });

        assert!(acc.usage().is_some());
        assert_eq!(acc.usage().map(|u| u.input_tokens), Some(100));
        assert!(matches!(acc.stop_reason(), Some(StopReason::ToolUse)));

        let blocks = acc.into_content_blocks();
        assert_eq!(blocks.len(), 2);
        assert!(matches!(&blocks[0], ContentBlock::Text { .. }));
        assert!(matches!(&blocks[1], ContentBlock::ToolUse { .. }));
    }

    #[test]
    fn test_accumulator_invalid_tool_json() {
        let mut acc = StreamAccumulator::new();

        acc.apply(&StreamDelta::ToolUseStart {
            id: "call_789".to_string(),
            name: "test_tool".to_string(),
            block_index: 0,
            thought_signature: None,
        });
        acc.apply(&StreamDelta::ToolInputDelta {
            id: "call_789".to_string(),
            delta: "invalid json {".to_string(),
            block_index: 0,
        });

        let blocks = acc.into_content_blocks();
        assert_eq!(blocks.len(), 1);
        match &blocks[0] {
            ContentBlock::ToolUse { input, .. } => {
                assert!(input.is_object());
            }
            _ => panic!("Expected ToolUse block"),
        }
    }

    #[test]
    fn test_accumulator_empty_tool_input_falls_back_to_empty_object() {
        // If no ToolInputDelta is received (e.g., stream interrupted or
        // deltas had mismatched IDs), the tool use block should still be
        // produced with an empty object so that the error is attributable
        // to the tool rather than silently lost.
        let mut acc = StreamAccumulator::new();

        acc.apply(&StreamDelta::ToolUseStart {
            id: "call_empty".to_string(),
            name: "read".to_string(),
            block_index: 0,
            thought_signature: None,
        });
        // No ToolInputDelta applied

        let blocks = acc.into_content_blocks();
        assert_eq!(blocks.len(), 1);
        match &blocks[0] {
            ContentBlock::ToolUse { input, name, .. } => {
                assert_eq!(name, "read");
                assert_eq!(input, &serde_json::json!({}));
            }
            _ => panic!("Expected ToolUse block"),
        }
    }

    #[test]
    fn test_accumulator_mismatched_delta_id_drops_input() {
        // If ToolInputDelta has a different ID than any ToolUseStart,
        // the input is silently dropped (the tool gets empty {}).
        let mut acc = StreamAccumulator::new();

        acc.apply(&StreamDelta::ToolUseStart {
            id: "call_A".to_string(),
            name: "bash".to_string(),
            block_index: 0,
            thought_signature: None,
        });
        // Delta with wrong ID
        acc.apply(&StreamDelta::ToolInputDelta {
            id: "call_B".to_string(),
            delta: r#"{"command":"ls"}"#.to_string(),
            block_index: 0,
        });

        let blocks = acc.into_content_blocks();
        assert_eq!(blocks.len(), 1);
        match &blocks[0] {
            ContentBlock::ToolUse { input, .. } => {
                // Input should be empty because the delta had a mismatched ID
                assert_eq!(input, &serde_json::json!({}));
            }
            _ => panic!("Expected ToolUse block"),
        }
    }

    #[test]
    fn test_accumulator_empty() {
        let acc = StreamAccumulator::new();
        let blocks = acc.into_content_blocks();
        assert!(blocks.is_empty());
    }

    #[test]
    fn test_accumulator_skips_empty_text() {
        let mut acc = StreamAccumulator::new();

        acc.apply(&StreamDelta::TextDelta {
            delta: String::new(),
            block_index: 0,
        });
        acc.apply(&StreamDelta::TextDelta {
            delta: "Hello".to_string(),
            block_index: 1,
        });

        let blocks = acc.into_content_blocks();
        assert_eq!(blocks.len(), 1);
        assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "Hello"));
    }
}