rig_harness/
progress.rs

1//! How far a running call got, for every front: [`ToolProgress`], and the
2//! output a running call's tool reports as it works ([`CallOutput`]).
3
4use bevy_app::prelude::*;
5use bevy_ecs::prelude::*;
6use crossbeam_channel::Receiver;
7use rig_core::tool::{ToolErrorKind, ToolExecutionError, ToolResult};
8use rig_ecs::stage::{AgentPass, IfStopped, Outcome, Stage};
9use rig_ecs::tasks::Wake;
10use rig_tools::{Live, OutputLimit};
11
12/// A line about a running call's progress, such as `12 of 40 files` or a
13/// command's last line of output, that the call's owner keeps current on
14/// the call's entity. The terminal view draws it under the call, and
15/// `--print` writes each change to stderr.
16#[derive(Component, Clone, Debug, Default, PartialEq, Eq)]
17pub struct ToolProgress(pub String);
18
19/// On a running call: the output its tool reports as it runs, through the
20/// [`Live`] it was made with. The harness keeps the call's
21/// [`ToolProgress`] (its last line) and its [`IfStopped`] (the output so
22/// far, cut as tool output is cut, then `[stopped before it finished]`)
23/// current from it, in [`Stage::Arrive`], so a call that is stopped answers
24/// with what it printed. The `shell` tool reports this way; a tool of your
25/// own does too:
26///
27/// ```
28/// use rig_harness::prelude::*;
29///
30/// /// Gives each `build` call a `Live` its tool writes its output to.
31/// fn follow_builds(
32///     calls: Query<(Entity, &ToolRequest), (ToolCalls, Added<Admitted>)>,
33///     wake: Res<Wake>,
34///     mut commands: Commands,
35/// ) {
36///     for (call, request) in &calls {
37///         if request.call.function.name.as_str() == "build" {
38///             let (output, _live) = CallOutput::new(&wake);
39///             commands.entity(call).insert(output);
40///             // The call's handler, built with `_live` in `PrepareSystems::Edit`,
41///             // calls `_live.line(..)` as it works.
42///         }
43///     }
44/// }
45/// ```
46#[derive(Component, Debug)]
47pub struct CallOutput {
48    /// What the tool reported and was not read yet.
49    lines: Receiver<String>,
50    /// The output read so far, at most about twice what the cut keeps.
51    kept: String,
52    /// Lines left out before `kept`.
53    cut_lines: usize,
54}
55
56impl CallOutput {
57    /// A call's output, and the [`Live`] its tool reports into; each piece
58    /// wakes the app.
59    pub fn new(wake: &Wake) -> (Self, Live) {
60        let (sender, lines) = crossbeam_channel::unbounded();
61        let wake = wake.clone();
62        let live = Live::new(move |text: &str| {
63            if sender.send(text.to_owned()).is_ok() {
64                wake.wake();
65            }
66        });
67        let output = Self {
68            lines,
69            kept: String::new(),
70            cut_lines: 0,
71        };
72        (output, live)
73    }
74
75    /// The output so far, cut as tool output is cut: its last lines, after
76    /// a line saying how many earlier ones were left out.
77    pub fn text(&self) -> String {
78        let (kept, dropped) = OutputLimit::default().tail(&self.kept);
79        match self.cut_lines + dropped {
80            0 => kept.to_owned(),
81            cut => format!("[{cut} earlier lines cut]\n{kept}"),
82        }
83    }
84
85    /// Reads what came since the last read; whether anything did.
86    fn read(&mut self) -> bool {
87        let mut came = false;
88        for piece in self.lines.try_iter() {
89            self.kept.push_str(&piece);
90            came = true;
91        }
92        let cut = OutputLimit::default();
93        if self.kept.len() > 2 * cut.bytes {
94            let (kept, dropped) = cut.tail(&self.kept);
95            self.cut_lines += dropped;
96            self.kept = kept.to_owned();
97        }
98        came
99    }
100
101    /// What the call answers if it is stopped now.
102    fn if_stopped(&self) -> IfStopped {
103        let mut text = self.text();
104        if !text.is_empty() && !text.ends_with('\n') {
105            text.push('\n');
106        }
107        text.push_str("[stopped before it finished]");
108        let stopped =
109            ToolExecutionError::new(ToolErrorKind::Cancelled, "stopped before it finished")
110                .with_model_feedback(text);
111        IfStopped(Outcome::from(ToolResult::failed(stopped)))
112    }
113}
114
115/// Follows each running call's [`CallOutput`].
116pub(crate) struct ProgressPlugin;
117
118impl Plugin for ProgressPlugin {
119    fn build(&self, app: &mut App) {
120        app.add_systems(AgentPass, follow_call_output.in_set(Stage::Arrive));
121    }
122}
123
124/// Reads each [`CallOutput`] and, when output came, writes its call's
125/// [`ToolProgress`] and [`IfStopped`].
126fn follow_call_output(
127    mut calls: Query<(Entity, &mut CallOutput), Without<Outcome>>,
128    mut commands: Commands,
129) {
130    for (call, mut output) in &mut calls {
131        if !output.read() {
132            continue;
133        }
134        let last = output
135            .kept
136            .lines()
137            .rev()
138            .find(|line| !line.trim().is_empty());
139        let progress = ToolProgress(last.unwrap_or_default().trim_end().to_owned());
140        let stopped = output.if_stopped();
141        commands.entity(call).try_insert((progress, stopped));
142    }
143}
144
145#[cfg(test)]
146mod tests;
rig_harness/progress.rs

rig_harness/
progress.rs
