apollo/agent/
streaming.rs

1//! Streaming responses — chunked output for long-running tasks
2//! Compatible with HTTP Server-Sent Events (SSE)
3
4use serde::Serialize;
5use tokio::sync::mpsc;
6
7/// Streaming output chunk
8#[derive(Debug, Clone, Serialize)]
9pub struct StreamChunk {
10    pub id: String,
11    pub chunk: String,
12    pub is_tool_use: bool,
13    pub tool_name: Option<String>,
14    pub index: usize,
15}
16
17/// Stream receiver — use for WebSocket/SSE responses
18pub type StreamReceiver = mpsc::UnboundedReceiver<StreamChunk>;
19
20/// Stream sender — used internally by agent loop
21pub type StreamSender = mpsc::UnboundedSender<StreamChunk>;
22
23/// Create a streaming channel
24pub fn stream_channel(_id: &str) -> (StreamSender, StreamReceiver) {
25    let (tx, rx) = mpsc::unbounded_channel();
26    (tx, rx)
27}
28
29/// Collect all chunks into a final response
30pub async fn collect_stream(mut rx: StreamReceiver) -> String {
31    let mut output = String::new();
32    while let Some(chunk) = rx.recv().await {
33        output.push_str(&chunk.chunk);
34    }
35    output
36}
37
38#[cfg(test)]
39mod tests {
40    use super::*;
41
42    #[tokio::test]
43    async fn test_stream_chunks() {
44        let (tx, rx) = stream_channel("test");
45
46        tokio::spawn(async move {
47            for i in 0..3 {
48                let chunk = StreamChunk {
49                    id: "test".to_string(),
50                    chunk: format!("chunk {}", i),
51                    is_tool_use: false,
52                    tool_name: None,
53                    index: i,
54                };
55                let _ = tx.send(chunk);
56            }
57        });
58
59        let output = collect_stream(rx).await;
60        assert_eq!(output, "chunk 0chunk 1chunk 2");
61    }
62}
apollo/agent/streaming.rs

apollo/agent/
streaming.rs
