agentos_client/
stream.rs

1//! Streaming / subscription primitives.
2//!
3//! Implements `spec.md` §5 / ADR-001 §5. The TypeScript `on*(id, handler) -> unsubscribe` pattern
4//! becomes streams + a uniform RAII [`Subscription`] guard:
5//!
6//! - process stdout/stderr, shell data, session events, permission requests, cron events ->
7//!   [`tokio::sync::broadcast`] (multi-subscriber; no replay).
8//! - process exit -> [`tokio::sync::watch`] seeded `None` (already-exited branch fires immediately
9//!   because the watch already holds `Some(code)`).
10//! - permission responder + internal single-reply correlation -> [`tokio::sync::oneshot`].
11
12use std::pin::Pin;
13use std::task::{Context, Poll};
14
15use futures::Stream;
16use tokio::sync::broadcast;
17use tokio_util::sync::ReusableBoxFuture;
18
19type ByteRecvResult = Result<Vec<u8>, broadcast::error::RecvError>;
20type ByteRecvState = (ByteRecvResult, broadcast::Receiver<Vec<u8>>);
21
22/// RAII guard returned by `on_*` register methods. Dropping it deregisters the subscription.
23///
24/// For broadcast/watch-backed subscriptions, dropping the returned stream/receiver is itself the
25/// unsubscribe; this guard wraps an optional deregistration closure for the cases (idempotent
26/// handler removal) that need explicit cleanup.
27#[must_use = "dropping the Subscription immediately unsubscribes"]
28pub struct Subscription {
29    on_drop: Option<Box<dyn FnOnce() + Send + Sync>>,
30}
31
32impl Subscription {
33    /// Create a subscription guard whose `Drop` runs `on_drop`.
34    pub fn new(on_drop: impl FnOnce() + Send + Sync + 'static) -> Self {
35        Self {
36            on_drop: Some(Box::new(on_drop)),
37        }
38    }
39
40    /// Create a no-op subscription guard (used when dropping the returned stream is the unsubscribe).
41    pub fn noop() -> Self {
42        Self { on_drop: None }
43    }
44
45    /// Detach the guard so dropping it no longer deregisters (subscription becomes permanent).
46    pub fn detach(mut self) {
47        self.on_drop = None;
48    }
49}
50
51impl std::fmt::Debug for Subscription {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        f.debug_struct("Subscription")
54            .field("active", &self.on_drop.is_some())
55            .finish()
56    }
57}
58
59impl Drop for Subscription {
60    fn drop(&mut self) {
61        if let Some(on_drop) = self.on_drop.take() {
62            on_drop();
63        }
64    }
65}
66
67/// A byte stream over a broadcast channel (process stdout/stderr, shell data).
68///
69/// Lagged messages are skipped. Closing the sender ends the stream.
70pub struct ByteStream {
71    inner: ReusableBoxFuture<'static, ByteRecvState>,
72}
73
74impl ByteStream {
75    /// Wrap a broadcast receiver as a [`Stream`] of byte chunks.
76    pub fn new(rx: broadcast::Receiver<Vec<u8>>) -> Self {
77        Self {
78            inner: ReusableBoxFuture::new(recv_bytes(rx)),
79        }
80    }
81}
82
83async fn recv_bytes(mut rx: broadcast::Receiver<Vec<u8>>) -> ByteRecvState {
84    let result = rx.recv().await;
85    (result, rx)
86}
87
88impl Stream for ByteStream {
89    type Item = Vec<u8>;
90
91    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
92        loop {
93            let (result, rx) = match self.inner.poll(cx) {
94                Poll::Ready(value) => value,
95                Poll::Pending => return Poll::Pending,
96            };
97            self.inner.set(recv_bytes(rx));
98            match result {
99                Ok(bytes) => return Poll::Ready(Some(bytes)),
100                Err(broadcast::error::RecvError::Lagged(_)) => continue,
101                Err(broadcast::error::RecvError::Closed) => return Poll::Ready(None),
102            }
103        }
104    }
105}
agentos_client/stream.rs

agentos_client/
stream.rs
