machi_tools/
stream.rs

1//! Streaming tool execution protocol.
2//!
3//! Invariant: a tool stream yields zero or more [`ToolStreamItem::Progress`]
4//! items followed by **exactly one** [`ToolStreamItem::Terminal`].
5
6use std::future::Future;
7use std::pin::Pin;
8
9use futures::Stream;
10use futures::stream::{self, StreamExt};
11use serde::{Deserialize, Serialize};
12use serde_json::Value;
13
14use crate::error::{ToolError, codes};
15use crate::tool::ToolResult;
16
17/// Opaque pinned stream of tool items.
18pub type ToolStream = Pin<Box<dyn Stream<Item = ToolStreamItem> + Send>>;
19
20/// One item in a tool stream.
21#[derive(Debug)]
22pub enum ToolStreamItem {
23    /// Intermediate progress (logs, partial stdout, custom payloads).
24    Progress(ToolProgress),
25    /// Terminal result — always last.
26    Terminal(Result<ToolResult, ToolError>),
27}
28
29impl ToolStreamItem {
30    /// Whether this is the terminal item.
31    #[must_use]
32    pub const fn is_terminal(&self) -> bool {
33        matches!(self, Self::Terminal(_))
34    }
35}
36
37/// Default max bytes per partial delta frame (16 KiB).
38pub const MAX_DELTA_BYTES: usize = 16 * 1024;
39/// Default max total frame/stream bytes (16 MiB).
40pub const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
41
42/// Progress payload shapes.
43#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
44#[serde(tag = "kind", rename_all = "snake_case")]
45#[non_exhaustive]
46pub enum ToolProgress {
47    /// Free-form text chunk.
48    Text {
49        /// Chunk body.
50        text: String,
51    },
52    /// Incremental partial output (W3.3).
53    Partial {
54        /// UTF-8-safe delta slice.
55        delta: String,
56        /// Total bytes emitted so far (including this delta).
57        total_bytes: u64,
58        /// Whether this delta was truncated to the frame cap.
59        truncated: bool,
60        /// Byte gap skipped since last partial (e.g. after truncation).
61        gap: u64,
62    },
63    /// Tool-defined progress.
64    Custom {
65        /// Stable producer discriminator.
66        subkind: String,
67        /// Arbitrary payload.
68        payload: Value,
69    },
70}
71
72impl ToolProgress {
73    /// Text progress helper.
74    #[must_use]
75    pub fn text(text: impl Into<String>) -> Self {
76        Self::Text { text: text.into() }
77    }
78
79    /// Partial progress helper.
80    #[must_use]
81    pub fn partial(delta: impl Into<String>, total_bytes: u64, truncated: bool, gap: u64) -> Self {
82        Self::Partial {
83            delta: delta.into(),
84            total_bytes,
85            truncated,
86            gap,
87        }
88    }
89}
90
91/// Split `input` into UTF-8-safe partial progress frames.
92///
93/// Each frame's `delta` is at most `max_delta_bytes` and ends on a char boundary.
94/// Stops once cumulative bytes would exceed `max_frame_bytes`.
95#[must_use]
96pub fn partial_progress_frames(
97    input: &str,
98    max_delta_bytes: usize,
99    max_frame_bytes: usize,
100) -> Vec<ToolProgress> {
101    let max_delta = max_delta_bytes.max(1);
102    let max_frame = max_frame_bytes.max(max_delta);
103    let mut out = Vec::new();
104    let mut offset = 0usize;
105    let mut total: u64 = 0;
106    let bytes = input.as_bytes();
107    while offset < bytes.len() {
108        if total >= u64::try_from(max_frame).unwrap_or(u64::MAX) {
109            break;
110        }
111        let remaining_frame =
112            max_frame.saturating_sub(usize::try_from(total).unwrap_or(usize::MAX));
113        let want = max_delta
114            .min(remaining_frame)
115            .min(bytes.len().saturating_sub(offset));
116        if want == 0 {
117            break;
118        }
119        let end = utf8_floor_end(input, offset, offset.saturating_add(want));
120        if end <= offset {
121            // Single multi-byte char larger than remaining budget — skip with gap.
122            let next = input[offset..]
123                .chars()
124                .next()
125                .map_or(offset.saturating_add(1), |c| {
126                    offset.saturating_add(c.len_utf8())
127                });
128            let gap = u64::try_from(next.saturating_sub(offset)).unwrap_or(0);
129            out.push(ToolProgress::partial(String::new(), total, true, gap));
130            offset = next;
131            continue;
132        }
133        let delta = input.get(offset..end).unwrap_or("").to_owned();
134        let delta_len = u64::try_from(delta.len()).unwrap_or(0);
135        total = total.saturating_add(delta_len);
136        let truncated = end - offset < want && end < bytes.len();
137        out.push(ToolProgress::partial(delta, total, truncated, 0));
138        offset = end;
139    }
140    out
141}
142
143/// Largest `end` in `(start, start+want]` that is a char boundary of `s`.
144fn utf8_floor_end(s: &str, start: usize, end: usize) -> usize {
145    let end = end.min(s.len());
146    if end <= start {
147        return start;
148    }
149    if s.is_char_boundary(end) {
150        return end;
151    }
152    let mut e = end;
153    while e > start && !s.is_char_boundary(e) {
154        e = e.saturating_sub(1);
155    }
156    e
157}
158
159/// Single-item terminal stream from a completed result.
160#[must_use]
161pub fn terminal_only(result: Result<ToolResult, ToolError>) -> ToolStream {
162    Box::pin(stream::once(
163        async move { ToolStreamItem::Terminal(result) },
164    ))
165}
166
167/// Progress items then a terminal future.
168pub fn with_progress<I, F, Fut>(progress: I, terminal: F) -> ToolStream
169where
170    I: IntoIterator<Item = ToolProgress> + Send + 'static,
171    F: FnOnce() -> Fut + Send + 'static,
172    Fut: Future<Output = Result<ToolResult, ToolError>> + Send + 'static,
173{
174    let items: Vec<ToolStreamItem> = progress.into_iter().map(ToolStreamItem::Progress).collect();
175    let prog = stream::iter(items);
176    let term = stream::once(async move { ToolStreamItem::Terminal(terminal().await) });
177    Box::pin(prog.chain(term))
178}
179
180/// Drain a stream to the terminal result, discarding progress.
181///
182/// # Errors
183///
184/// Returns stream protocol error when the stream ends without a terminal item.
185pub async fn drain_terminal(stream: ToolStream) -> Result<ToolResult, ToolError> {
186    let (_progress, result) = drain_with_progress(stream).await;
187    result
188}
189
190/// Drain a stream, collecting progress items and the terminal result.
191///
192/// # Errors
193///
194/// The terminal `Result` carries tool failures. When the stream ends without a
195/// terminal item, returns protocol error in the terminal slot.
196pub async fn drain_with_progress(
197    mut stream: ToolStream,
198) -> (Vec<ToolProgress>, Result<ToolResult, ToolError>) {
199    let mut progress = Vec::new();
200    while let Some(item) = stream.next().await {
201        match item {
202            ToolStreamItem::Progress(p) => progress.push(p),
203            ToolStreamItem::Terminal(result) => return (progress, result),
204        }
205    }
206    (
207        progress,
208        Err(codes::stream_protocol(
209            "tool stream ended without a terminal item",
210        )),
211    )
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217
218    #[tokio::test]
219    async fn terminal_only_drains() {
220        let s = terminal_only(Ok(ToolResult::text("ok")));
221        let r = drain_terminal(s).await.expect("drain");
222        assert_eq!(r.content, "ok");
223    }
224
225    #[tokio::test]
226    async fn progress_then_terminal() {
227        let s = with_progress(vec![ToolProgress::text("working")], || async {
228            Ok(ToolResult::text("done"))
229        });
230        let r = drain_terminal(s).await.expect("drain");
231        assert_eq!(r.content, "done");
232    }
233
234    #[tokio::test]
235    async fn drain_with_progress_captures_chunks() {
236        let s = with_progress(
237            vec![ToolProgress::text("a"), ToolProgress::text("b")],
238            || async { Ok(ToolResult::text("done")) },
239        );
240        let (progress, result) = drain_with_progress(s).await;
241        assert_eq!(progress.len(), 2);
242        assert_eq!(result.expect("ok").content, "done");
243    }
244
245    #[tokio::test]
246    async fn empty_stream_is_protocol_error() {
247        let s: ToolStream = Box::pin(stream::empty());
248        let err = drain_terminal(s).await.expect_err("proto");
249        assert_eq!(err.code(), machi_types::ErrorCode::ToolStreamProtocol);
250    }
251
252    #[test]
253    fn partial_frames_respect_utf8_and_caps() {
254        let s = "hello🎉world";
255        let frames = partial_progress_frames(s, 4, 10_000);
256        assert!(!frames.is_empty());
257        let mut rebuilt = String::new();
258        for f in &frames {
259            if let ToolProgress::Partial { delta, .. } = f {
260                rebuilt.push_str(delta);
261            }
262        }
263        // May skip oversized multi-byte with gap; all deltas must be valid UTF-8.
264        assert!(rebuilt.is_char_boundary(rebuilt.len()));
265        for f in frames {
266            if let ToolProgress::Partial { delta, .. } = f {
267                assert!(delta.len() <= 4 || delta.is_empty());
268            }
269        }
270    }
271
272    #[test]
273    fn partial_frames_honor_frame_budget() {
274        let s = "abcdefghij";
275        let frames = partial_progress_frames(s, 3, 6);
276        let total: usize = frames
277            .iter()
278            .map(|f| match f {
279                ToolProgress::Partial { delta, .. } => delta.len(),
280                _ => 0,
281            })
282            .sum();
283        assert!(total <= 6);
284    }
285}
machi_tools/stream.rs

machi_tools/
stream.rs
