aither_core/llm/tool/
context.rs

1//! The per-call context a [`Tool`](super::Tool) runs with.
2//!
3//! Every invocation of [`Tool::call`](super::Tool::call) receives its own
4//! [`ToolContext`]. Whoever dispatches the call decides what the context is
5//! connected to: an MCP server connects it to the request's progress token, an
6//! agent that has nowhere to show progress passes [`ToolContext::new`]. The
7//! tool reports the same way in both cases and never learns which it got.
8
9use alloc::boxed::Box;
10use alloc::string::String;
11use core::fmt::{self, Debug, Display};
12use core::future::Future;
13use core::pin::Pin;
14
15/// How far a running tool call has got.
16///
17/// `progress` is the amount of work done so far and `total`, when known, the
18/// amount there is; both may be fractional and neither has a unit. `message`
19/// is a human-readable line describing the current state.
20#[derive(Debug, Clone, PartialEq)]
21pub struct Progress {
22    done: f64,
23    total: Option<f64>,
24    message: Option<String>,
25}
26
27impl Progress {
28    /// Progress with the given amount of work done, no total, and no message.
29    #[must_use]
30    pub const fn new(progress: f64) -> Self {
31        Self {
32            done: progress,
33            total: None,
34            message: None,
35        }
36    }
37
38    /// Sets the total amount of work.
39    #[must_use]
40    pub const fn with_total(mut self, total: f64) -> Self {
41        self.total = Some(total);
42        self
43    }
44
45    /// Sets a human-readable description of the current state.
46    #[must_use]
47    pub fn with_message(mut self, message: impl Into<String>) -> Self {
48        self.message = Some(message.into());
49        self
50    }
51
52    /// The amount of work done so far.
53    #[must_use]
54    pub const fn progress(&self) -> f64 {
55        self.done
56    }
57
58    /// The total amount of work, if known.
59    #[must_use]
60    pub const fn total(&self) -> Option<f64> {
61        self.total
62    }
63
64    /// The human-readable description of the current state, if any.
65    #[must_use]
66    pub fn message(&self) -> Option<&str> {
67        self.message.as_deref()
68    }
69}
70
71/// Where a [`ToolContext`] delivers the progress its tool reports.
72///
73/// Implemented by whatever dispatches tool calls and has somewhere to show
74/// progress, such as an MCP server forwarding it to the client as
75/// `notifications/progress`. Every report reaching a sink has already been
76/// validated by [`ToolContext::report_progress`]: its values are finite and
77/// its `progress` is greater than that of the report before it.
78pub trait ProgressSink: Send + 'static {
79    /// Delivers one progress report.
80    ///
81    /// The returned future may wait for room to deliver (backpressure) but
82    /// must not block the thread. A sink whose destination is gone drops the
83    /// report: progress is advisory, and the call it belongs to is being torn
84    /// down anyway.
85    fn report(&mut self, progress: Progress) -> impl Future<Output = ()> + Send;
86}
87
88/// Object-safe twin of [`ProgressSink`], so [`ToolContext`] can hold any sink.
89trait ProgressSinkImpl: Send {
90    fn report(&mut self, progress: Progress) -> Pin<Box<dyn Future<Output = ()> + Send + '_>>;
91}
92
93impl<S: ProgressSink> ProgressSinkImpl for S {
94    fn report(&mut self, progress: Progress) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
95        Box::pin(ProgressSink::report(self, progress))
96    }
97}
98
99/// A progress report that no receiver could make sense of.
100///
101/// Both variants are bugs in the reporting tool, so they are returned to it
102/// rather than papered over.
103#[derive(Debug, Clone, PartialEq)]
104pub enum ProgressError {
105    /// `progress` or `total` was NaN or infinite.
106    NotFinite(Progress),
107    /// `progress` did not exceed the previous report's. Progress must increase
108    /// with every report, even when the total is unknown.
109    NotIncreasing {
110        /// The `progress` of the previous accepted report.
111        previous: f64,
112        /// The rejected report.
113        reported: Progress,
114    },
115}
116
117impl Display for ProgressError {
118    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
119        match self {
120            Self::NotFinite(progress) => write!(
121                f,
122                "progress report is not finite: progress {}, total {:?}",
123                progress.done, progress.total
124            ),
125            Self::NotIncreasing { previous, reported } => write!(
126                f,
127                "progress must increase with every report: {} does not exceed {previous}",
128                reported.done
129            ),
130        }
131    }
132}
133
134impl core::error::Error for ProgressError {}
135
136/// The context of one tool call.
137///
138/// A tool receives it by value in [`Tool::call`](super::Tool::call) and may
139/// move it wherever the work happens. It is deliberately not `Clone`: one
140/// owner reports in sequence, which is what lets it guarantee that progress
141/// increases report by report.
142pub struct ToolContext {
143    sink: Option<Box<dyn ProgressSinkImpl>>,
144    last_progress: Option<f64>,
145}
146
147impl ToolContext {
148    /// A context whose caller does not listen for progress.
149    ///
150    /// Reports are still validated, so a tool that misreports fails the same
151    /// way whether or not anyone is listening, but they are delivered nowhere.
152    #[must_use]
153    pub const fn new() -> Self {
154        Self {
155            sink: None,
156            last_progress: None,
157        }
158    }
159
160    /// A context that delivers its tool's progress reports to `sink`.
161    #[must_use]
162    pub fn with_progress(sink: impl ProgressSink) -> Self {
163        Self {
164            sink: Some(Box::new(sink)),
165            last_progress: None,
166        }
167    }
168
169    /// Whether the caller listens for progress.
170    ///
171    /// A tool whose work can outlast its caller's idle limit decides from
172    /// this whether reporting progress keeps the call alive: without a
173    /// listener, reports go nowhere and the caller sees nothing until the
174    /// call returns.
175    #[must_use]
176    pub const fn is_listening(&self) -> bool {
177        self.sink.is_some()
178    }
179
180    /// Reports how far the call has got.
181    ///
182    /// Delivered to the caller when it listens for progress, and a no-op when
183    /// it does not. Awaiting the returned future may wait for the caller to
184    /// take earlier reports.
185    ///
186    /// # Errors
187    ///
188    /// Returns [`ProgressError::NotFinite`] when `progress` or `total` is NaN
189    /// or infinite, and [`ProgressError::NotIncreasing`] when `progress` does
190    /// not exceed the previous report's. A rejected report is not delivered
191    /// and does not count as the previous report.
192    pub async fn report_progress(&mut self, progress: Progress) -> Result<(), ProgressError> {
193        if !progress.done.is_finite() || progress.total.is_some_and(|total| !total.is_finite()) {
194            return Err(ProgressError::NotFinite(progress));
195        }
196        if let Some(previous) = self.last_progress
197            && progress.done <= previous
198        {
199            return Err(ProgressError::NotIncreasing {
200                previous,
201                reported: progress,
202            });
203        }
204        self.last_progress = Some(progress.done);
205        if let Some(sink) = &mut self.sink {
206            sink.report(progress).await;
207        }
208        Ok(())
209    }
210}
211
212impl Default for ToolContext {
213    fn default() -> Self {
214        Self::new()
215    }
216}
217
218impl Debug for ToolContext {
219    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
220        f.debug_struct("ToolContext")
221            .field("listening", &self.sink.is_some())
222            .field("last_progress", &self.last_progress)
223            .finish()
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    extern crate std;
230
231    use super::*;
232    use alloc::vec::Vec;
233    use std::sync::mpsc;
234
235    /// Forwards every report it is handed to a channel the test drains.
236    struct Collect(mpsc::Sender<Progress>);
237
238    impl ProgressSink for Collect {
239        async fn report(&mut self, progress: Progress) {
240            self.0.send(progress).expect("test receiver alive");
241        }
242    }
243
244    fn listening() -> (ToolContext, mpsc::Receiver<Progress>) {
245        let (tx, rx) = mpsc::channel();
246        (ToolContext::with_progress(Collect(tx)), rx)
247    }
248
249    #[test]
250    fn reports_whether_the_caller_listens() {
251        assert!(!ToolContext::new().is_listening());
252        assert!(listening().0.is_listening());
253    }
254
255    #[tokio::test]
256    async fn delivers_increasing_reports_in_order() {
257        let (mut cx, seen) = listening();
258
259        cx.report_progress(Progress::new(1.0)).await.unwrap();
260        cx.report_progress(Progress::new(2.5).with_total(10.0).with_message("half"))
261            .await
262            .unwrap();
263
264        assert_eq!(
265            seen.try_iter().collect::<Vec<_>>(),
266            [
267                Progress::new(1.0),
268                Progress::new(2.5).with_total(10.0).with_message("half"),
269            ]
270        );
271    }
272
273    #[tokio::test]
274    async fn rejects_non_increasing_and_non_finite_reports() {
275        let (mut cx, seen) = listening();
276
277        cx.report_progress(Progress::new(3.0)).await.unwrap();
278        assert_eq!(
279            cx.report_progress(Progress::new(3.0)).await,
280            Err(ProgressError::NotIncreasing {
281                previous: 3.0,
282                reported: Progress::new(3.0),
283            })
284        );
285        assert!(matches!(
286            cx.report_progress(Progress::new(f64::NAN)).await,
287            Err(ProgressError::NotFinite(_))
288        ));
289        assert!(matches!(
290            cx.report_progress(Progress::new(4.0).with_total(f64::INFINITY))
291                .await,
292            Err(ProgressError::NotFinite(_))
293        ));
294        // A rejected report is not the new baseline.
295        cx.report_progress(Progress::new(3.5)).await.unwrap();
296
297        assert_eq!(
298            seen.try_iter().collect::<Vec<_>>(),
299            [Progress::new(3.0), Progress::new(3.5)]
300        );
301    }
302
303    #[tokio::test]
304    async fn context_without_a_listener_still_validates() {
305        let mut cx = ToolContext::new();
306        cx.report_progress(Progress::new(1.0)).await.unwrap();
307        assert!(cx.report_progress(Progress::new(0.5)).await.is_err());
308    }
309}
aither_core/llm/tool/context.rs

aither_core/llm/tool/
context.rs
