antigravity_codes/
steps.rs

1//! Turning a stream of [`StepUpdate`] frames into coherent steps.
2//!
3//! The harness reports a step many times as it runs: first as `STATE_ACTIVE`
4//! with `text_delta`/`thinking_delta` fragments, then once more as
5//! `STATE_DONE` carrying the whole `text`. A step is identified by
6//! `(trajectory_id, step_index)` — *not* by index alone, because subagents run
7//! on their own trajectories concurrently with the main one.
8//!
9//! [`StepAssembler`] keeps the running text for each of those keys and hands
10//! back a [`Step`] snapshot per update, so a caller can render incrementally
11//! and still see a complete step at the end.
12
13use std::collections::HashMap;
14
15use crate::protocol::{StepUpdate, StepUpdateSource, StepUpdateState, StepUpdateTarget};
16
17/// What a step was doing.
18///
19/// Derived from whichever action member the harness populated. Steps that
20/// carry only prose land on [`StepKind::Message`].
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
22#[non_exhaustive]
23pub enum StepKind {
24    /// Prose for the user, or the echo of their own input.
25    Message,
26    /// A directory listing.
27    ListDirectory,
28    /// A filename search.
29    FindFile,
30    /// A content search across a directory.
31    SearchDirectory,
32    /// A file read.
33    ViewFile,
34    /// A file creation.
35    CreateFile,
36    /// A file edit, with the diff attached.
37    EditFile,
38    /// A shell command, with its exit code and combined output.
39    RunCommand,
40    /// The harness compacting its own context.
41    Compaction,
42    /// Delegation to a subagent, which runs on its own trajectory.
43    InvokeSubagent,
44    /// Image generation.
45    GenerateImage,
46    /// A web search.
47    SearchWeb,
48    /// A URL fetch.
49    ReadUrlContent,
50    /// A tool call routed to an MCP server.
51    McpTool,
52    /// A tool the *client* is expected to run.
53    CustomTool,
54    /// The agent's closing summary.
55    Finish,
56    /// A failure the harness is reporting in-band.
57    Error,
58    /// The harness is waiting for the user to confirm a tool.
59    ToolConfirmationRequest,
60    /// The harness is asking the user a question.
61    QuestionsRequest,
62}
63
64/// A step, merged across every update the harness has sent for it.
65#[derive(Debug, Clone)]
66pub struct Step {
67    /// The trajectory this step belongs to. The main conversation and each
68    /// subagent get their own.
69    pub trajectory_id: String,
70    /// Position within the trajectory.
71    pub step_index: u32,
72    /// Lifecycle state as of this update.
73    pub state: StepUpdateState,
74    /// Who produced the step.
75    pub source: StepUpdateSource,
76    /// Who it is addressed to.
77    pub target: StepUpdateTarget,
78    /// What the step is doing.
79    pub kind: StepKind,
80    /// Text accumulated so far, deltas included.
81    pub text: String,
82    /// Reasoning accumulated so far, deltas included.
83    pub thinking: String,
84    /// Only the text that arrived in *this* update, for incremental rendering.
85    pub text_delta: String,
86    /// Failure detail, when [`Self::state`] is [`StepUpdateState::Error`].
87    pub error_message: Option<String>,
88    /// The frame this snapshot came from, for anything the summary omits.
89    pub update: StepUpdate,
90}
91
92impl Step {
93    /// A stable key for this step: `"{trajectory_id}:{step_index}"`.
94    pub fn id(&self) -> String {
95        format!("{}:{}", self.trajectory_id, self.step_index)
96    }
97
98    /// True once the harness will not update this step again.
99    pub fn is_final(&self) -> bool {
100        matches!(self.state, StepUpdateState::Done | StepUpdateState::Error)
101    }
102
103    /// The accumulated text, if there is any.
104    pub fn text(&self) -> Option<&str> {
105        Some(self.text.as_str()).filter(|t| !t.is_empty())
106    }
107
108    /// Text addressed to the user, as opposed to the model or the environment.
109    ///
110    /// This is the filter to use when rendering a conversation: it drops the
111    /// echo of the user's own input and the agent's tool chatter.
112    pub fn user_facing_text(&self) -> Option<&str> {
113        match self.target {
114            StepUpdateTarget::User => self.text(),
115            _ => None,
116        }
117    }
118}
119
120/// Accumulates [`StepUpdate`] frames into [`Step`]s.
121#[derive(Debug, Default)]
122pub struct StepAssembler {
123    buffers: HashMap<(String, u32), Buffer>,
124    main_trajectory: Option<String>,
125}
126
127#[derive(Debug, Default)]
128struct Buffer {
129    text: String,
130    thinking: String,
131}
132
133impl StepAssembler {
134    /// A fresh assembler for a conversation.
135    ///
136    /// `main_trajectory` should be the conversation's cascade id when known;
137    /// otherwise the first trajectory seen is adopted as the main one.
138    pub fn new(main_trajectory: Option<String>) -> Self {
139        Self {
140            buffers: HashMap::new(),
141            main_trajectory,
142        }
143    }
144
145    /// The trajectory treated as the main conversation.
146    pub fn main_trajectory(&self) -> Option<&str> {
147        self.main_trajectory.as_deref()
148    }
149
150    /// True when `trajectory_id` is the main conversation rather than a
151    /// subagent's.
152    pub fn is_main(&self, trajectory_id: &str) -> bool {
153        self.main_trajectory.as_deref() == Some(trajectory_id)
154    }
155
156    /// Folds one update into the running state and returns the merged step.
157    pub fn ingest(&mut self, update: StepUpdate) -> Step {
158        let trajectory_id = update.trajectory_id.clone().unwrap_or_default();
159        let step_index = update.step_index.unwrap_or_default();
160        if self.main_trajectory.is_none() && !trajectory_id.is_empty() {
161            self.main_trajectory = Some(trajectory_id.clone());
162        }
163
164        let buffer = self
165            .buffers
166            .entry((trajectory_id.clone(), step_index))
167            .or_default();
168
169        let text_delta = update.text_delta.clone().unwrap_or_default();
170        if !text_delta.is_empty() {
171            buffer.text.push_str(&text_delta);
172        }
173        // The harness sends the whole text again when the step settles; that
174        // copy wins, since it is what the model actually committed to.
175        if let Some(text) = update.text.as_deref().filter(|t| !t.is_empty()) {
176            buffer.text = text.to_string();
177        }
178
179        if let Some(delta) = update.thinking_delta.as_deref().filter(|t| !t.is_empty()) {
180            buffer.thinking.push_str(delta);
181        }
182        if let Some(thinking) = update.thinking.as_deref().filter(|t| !t.is_empty()) {
183            buffer.thinking = thinking.to_string();
184        }
185
186        let step = Step {
187            trajectory_id,
188            step_index,
189            state: update.state.clone().unwrap_or_default(),
190            source: update.source.clone().unwrap_or_default(),
191            target: update.target.clone().unwrap_or_default(),
192            kind: classify(&update),
193            text: buffer.text.clone(),
194            thinking: buffer.thinking.clone(),
195            text_delta,
196            error_message: update.error_message.clone().filter(|m| !m.is_empty()),
197            update,
198        };
199
200        if step.is_final() {
201            self.buffers
202                .remove(&(step.trajectory_id.clone(), step.step_index));
203        }
204        step
205    }
206}
207
208fn classify(update: &StepUpdate) -> StepKind {
209    // Ordered so the more specific requests win over the action that carries
210    // them; a step waiting on a confirmation also has its action populated.
211    if update.tool_confirmation_request.is_some() {
212        StepKind::ToolConfirmationRequest
213    } else if update.questions_request.is_some() {
214        StepKind::QuestionsRequest
215    } else if update.error.is_some() {
216        StepKind::Error
217    } else if update.finish.is_some() {
218        StepKind::Finish
219    } else if update.list_directory.is_some() {
220        StepKind::ListDirectory
221    } else if update.find_file.is_some() {
222        StepKind::FindFile
223    } else if update.search_directory.is_some() {
224        StepKind::SearchDirectory
225    } else if update.view_file.is_some() {
226        StepKind::ViewFile
227    } else if update.create_file.is_some() {
228        StepKind::CreateFile
229    } else if update.edit_file.is_some() {
230        StepKind::EditFile
231    } else if update.run_command.is_some() {
232        StepKind::RunCommand
233    } else if update.compaction.is_some() {
234        StepKind::Compaction
235    } else if update.invoke_subagent.is_some() {
236        StepKind::InvokeSubagent
237    } else if update.generate_image.is_some() {
238        StepKind::GenerateImage
239    } else if update.search_web.is_some() {
240        StepKind::SearchWeb
241    } else if update.read_url_content.is_some() {
242        StepKind::ReadUrlContent
243    } else if update.mcp_tool.is_some() {
244        StepKind::McpTool
245    } else if update.custom_tool.is_some() {
246        StepKind::CustomTool
247    } else {
248        StepKind::Message
249    }
250}
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255    use crate::protocol::{ActionRunCommand, UserQuestionsRequest};
256
257    fn delta(trajectory: &str, index: u32, text: &str, state: StepUpdateState) -> StepUpdate {
258        StepUpdate {
259            trajectory_id: Some(trajectory.into()),
260            step_index: Some(index),
261            state: Some(state),
262            text_delta: Some(text.into()),
263            ..Default::default()
264        }
265    }
266
267    #[test]
268    fn deltas_accumulate_within_a_step() {
269        let mut a = StepAssembler::new(None);
270        a.ingest(delta("t1", 0, "Hel", StepUpdateState::Active));
271        let step = a.ingest(delta("t1", 0, "lo", StepUpdateState::Active));
272        assert_eq!(step.text, "Hello");
273        assert_eq!(step.text_delta, "lo");
274        assert!(!step.is_final());
275    }
276
277    #[test]
278    fn the_final_text_replaces_the_accumulated_deltas() {
279        let mut a = StepAssembler::new(None);
280        a.ingest(delta("t1", 0, "Hel", StepUpdateState::Active));
281        let step = a.ingest(StepUpdate {
282            trajectory_id: Some("t1".into()),
283            step_index: Some(0),
284            state: Some(StepUpdateState::Done),
285            text: Some("Hello, world".into()),
286            ..Default::default()
287        });
288        assert_eq!(step.text, "Hello, world");
289        assert!(step.is_final());
290    }
291
292    #[test]
293    fn concurrent_trajectories_do_not_bleed_into_each_other() {
294        let mut a = StepAssembler::new(Some("main".into()));
295        a.ingest(delta("main", 0, "main-", StepUpdateState::Active));
296        a.ingest(delta("sub", 0, "sub-", StepUpdateState::Active));
297        let main = a.ingest(delta("main", 0, "text", StepUpdateState::Active));
298        let sub = a.ingest(delta("sub", 0, "text", StepUpdateState::Active));
299        assert_eq!(main.text, "main-text");
300        assert_eq!(sub.text, "sub-text");
301        assert!(a.is_main("main"));
302        assert!(!a.is_main("sub"));
303    }
304
305    #[test]
306    fn the_first_trajectory_seen_becomes_the_main_one() {
307        let mut a = StepAssembler::new(None);
308        a.ingest(delta("first", 0, "x", StepUpdateState::Active));
309        assert_eq!(a.main_trajectory(), Some("first"));
310        a.ingest(delta("second", 0, "y", StepUpdateState::Active));
311        assert_eq!(a.main_trajectory(), Some("first"));
312    }
313
314    #[test]
315    fn a_settled_step_releases_its_buffer() {
316        let mut a = StepAssembler::new(None);
317        a.ingest(delta("t1", 0, "hi", StepUpdateState::Done));
318        assert!(a.buffers.is_empty());
319    }
320
321    #[test]
322    fn actions_classify_by_their_populated_member() {
323        let update = StepUpdate {
324            run_command: Some(ActionRunCommand {
325                command_line: Some("ls".into()),
326                ..Default::default()
327            }),
328            ..Default::default()
329        };
330        assert_eq!(classify(&update), StepKind::RunCommand);
331        assert_eq!(classify(&StepUpdate::default()), StepKind::Message);
332    }
333
334    #[test]
335    fn a_pending_question_outranks_its_action() {
336        let update = StepUpdate {
337            run_command: Some(ActionRunCommand::default()),
338            questions_request: Some(UserQuestionsRequest::default()),
339            ..Default::default()
340        };
341        assert_eq!(classify(&update), StepKind::QuestionsRequest);
342    }
343
344    #[test]
345    fn only_user_targeted_text_is_user_facing() {
346        let mut a = StepAssembler::new(None);
347        let to_model = a.ingest(StepUpdate {
348            trajectory_id: Some("t".into()),
349            target: Some(StepUpdateTarget::Model),
350            text: Some("echo of the prompt".into()),
351            ..Default::default()
352        });
353        assert_eq!(to_model.user_facing_text(), None);
354        assert_eq!(to_model.text(), Some("echo of the prompt"));
355    }
356}
antigravity_codes/steps.rs

antigravity_codes/
steps.rs
