apollo/agent/
stream.rs

1use serde::Serialize;
2use tokio::sync::mpsc;
3
4#[derive(Debug, Clone, Serialize)]
5#[serde(tag = "type", rename_all = "snake_case")]
6pub enum AgentStreamEvent {
7    Status {
8        message: String,
9    },
10    ToolStart {
11        name: String,
12        hint: String,
13    },
14    ToolEnd {
15        name: String,
16        ok: bool,
17        elapsed_secs: u64,
18    },
19    Delta {
20        text: String,
21    },
22    Done {
23        response: String,
24    },
25    Error {
26        message: String,
27    },
28}
29
30pub type AgentStreamTx = mpsc::UnboundedSender<AgentStreamEvent>;
31
32pub fn emit(tx: &Option<AgentStreamTx>, event: AgentStreamEvent) {
33    if let Some(tx) = tx {
34        let _ = tx.send(event);
35    }
36}
37
38tokio::task_local! {
39    /// Sink for the turn running in the current task.
40    ///
41    /// A single global sink on `AgentRunner` cannot serve two clients: the
42    /// second connection overwrites the first one's sender, so A's tokens go
43    /// to B and A's turn end clears B's sink. Scoping the sink to the task
44    /// that runs the turn keeps concurrent turns independent.
45    static TURN_STREAM: Option<AgentStreamTx>;
46}
47
48/// The sink for the turn running in this task, if any.
49pub fn current_turn_sink() -> Option<AgentStreamTx> {
50    TURN_STREAM.try_with(|tx| tx.clone()).ok().flatten()
51}
52
53/// Run `fut` with `tx` as the sink for this turn.
54pub async fn with_turn_sink<F>(tx: Option<AgentStreamTx>, fut: F) -> F::Output
55where
56    F: std::future::Future,
57{
58    TURN_STREAM.scope(tx, fut).await
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64
65    async fn one_turn(label: &str, tx: AgentStreamTx) {
66        with_turn_sink(Some(tx), async {
67            for i in 0..50 {
68                emit(
69                    &current_turn_sink(),
70                    AgentStreamEvent::Delta {
71                        text: format!("{label}{i}"),
72                    },
73                );
74                tokio::task::yield_now().await;
75            }
76        })
77        .await;
78    }
79
80    #[tokio::test]
81    async fn concurrent_turns_do_not_cross_streams() {
82        let (tx_a, mut rx_a) = mpsc::unbounded_channel();
83        let (tx_b, mut rx_b) = mpsc::unbounded_channel();
84
85        let a = tokio::spawn(one_turn("a", tx_a));
86        let b = tokio::spawn(one_turn("b", tx_b));
87        a.await.unwrap();
88        b.await.unwrap();
89
90        let drain = |rx: &mut mpsc::UnboundedReceiver<AgentStreamEvent>| {
91            let mut out = Vec::new();
92            while let Ok(AgentStreamEvent::Delta { text }) = rx.try_recv() {
93                out.push(text);
94            }
95            out
96        };
97
98        let a_events = drain(&mut rx_a);
99        let b_events = drain(&mut rx_b);
100        assert_eq!(a_events.len(), 50);
101        assert_eq!(b_events.len(), 50);
102        assert!(a_events.iter().all(|t| t.starts_with('a')));
103        assert!(b_events.iter().all(|t| t.starts_with('b')));
104    }
105
106    #[tokio::test]
107    async fn no_scope_means_no_turn_sink() {
108        assert!(current_turn_sink().is_none());
109    }
110}
apollo/agent/stream.rs

apollo/agent/
stream.rs
