Skip to main content

wisp/components/
tool_call_statuses.rs

1use acp_utils::notifications::SubAgentProgressParams;
2use agent_client_protocol as acp;
3use std::collections::HashMap;
4use std::time::Instant;
5
6use crate::components::sub_agent_tracker::SubAgentTracker;
7use crate::components::tool_call_status_view::{
8    ToolCallStatus, compute_diff_preview, render_tool_tree,
9};
10use crate::components::tracked_tool_call::{
11    TrackedToolCall, raw_input_fragment, upsert_tracked_tool_call,
12};
13use tui::{Line, ViewContext};
14
15/// Tracks active tool calls and produces status lines for the frame.
16#[derive(Clone)]
17pub struct ToolCallStatuses {
18    /// Ordered list of tool call IDs (insertion order)
19    tool_order: Vec<String>,
20    /// Tool call info by ID
21    tool_calls: HashMap<String, TrackedToolCall>,
22    /// Sub-agent states keyed by parent tool call ID
23    sub_agents: SubAgentTracker,
24    /// Animation tick for the spinner on running tool calls
25    tick: u16,
26}
27
28pub struct ToolProgress {
29    pub running_any: bool,
30    pub completed_top_level: usize,
31    pub total_top_level: usize,
32}
33
34impl ToolCallStatuses {
35    pub fn new() -> Self {
36        Self {
37            tool_order: Vec::new(),
38            tool_calls: HashMap::new(),
39            sub_agents: SubAgentTracker::default(),
40            tick: 0,
41        }
42    }
43
44    pub fn progress(&self) -> ToolProgress {
45        let running_any = self.any_running_including_subagents();
46        let (completed_top_level, total_top_level) = self.top_level_counts();
47        ToolProgress {
48            running_any,
49            completed_top_level,
50            total_top_level,
51        }
52    }
53
54    /// Advance the animation state. Call this on tick events.
55    pub fn on_tick(&mut self, _now: Instant) {
56        if self.progress().running_any {
57            self.tick = self.tick.wrapping_add(1);
58        }
59    }
60
61    /// Handle a new tool call from ACP `SessionUpdate::ToolCall`.
62    pub fn on_tool_call(&mut self, tool_call: &acp::ToolCall) {
63        let id = tool_call.tool_call_id.0.to_string();
64        let arguments = tool_call
65            .raw_input
66            .as_ref()
67            .map(raw_input_fragment)
68            .unwrap_or_default();
69
70        let tracked = upsert_tracked_tool_call(
71            &mut self.tool_order,
72            &mut self.tool_calls,
73            &id,
74            &tool_call.title,
75            arguments.clone(),
76        );
77        tracked.update_name(&tool_call.title);
78        tracked.arguments = arguments;
79        tracked.status = ToolCallStatus::Running;
80    }
81
82    /// Handle a tool call update from ACP `SessionUpdate::ToolCallUpdate`.
83    pub fn on_tool_call_update(&mut self, update: &acp::ToolCallUpdate) {
84        let id = update.tool_call_id.0.to_string();
85
86        if let Some(tc) = self.tool_calls.get_mut(&id) {
87            if let Some(title) = &update.fields.title {
88                tc.update_name(title);
89            }
90            if let Some(raw_input) = &update.fields.raw_input {
91                tc.append_arguments(&raw_input_fragment(raw_input));
92            }
93            if let Some(meta) = &update.meta
94                && let Some(dv) = meta.get("display_value").and_then(|v| v.as_str())
95            {
96                tc.display_value = Some(dv.to_string());
97            }
98            if let Some(content) = &update.fields.content {
99                for item in content {
100                    if let acp::ToolCallContent::Diff(diff) = item {
101                        tc.diff_preview = Some(compute_diff_preview(diff));
102                    }
103                }
104            }
105            if let Some(status) = update.fields.status {
106                tc.apply_status(status);
107            }
108        }
109    }
110
111    pub fn finalize_running(&mut self, cancelled: bool) {
112        let terminal_status = if cancelled {
113            ToolCallStatus::Error("cancelled".to_string())
114        } else {
115            ToolCallStatus::Success
116        };
117
118        for tool_call in self.tool_calls.values_mut() {
119            if matches!(tool_call.status, ToolCallStatus::Running) {
120                tool_call.status = terminal_status.clone();
121            }
122        }
123
124        self.sub_agents.finalize_running(cancelled);
125    }
126
127    pub fn has_tool(&self, id: &str) -> bool {
128        self.tool_calls.contains_key(id)
129    }
130
131    #[cfg(test)]
132    pub fn is_tool_running(&self, id: &str) -> bool {
133        self.tool_calls
134            .get(id)
135            .is_some_and(|tc| matches!(tc.status, ToolCallStatus::Running))
136    }
137
138    /// Handle a sub-agent progress notification.
139    pub fn on_sub_agent_progress(&mut self, notification: &SubAgentProgressParams) {
140        self.sub_agents.on_progress(notification);
141    }
142
143    #[cfg(test)]
144    pub fn remove_tool(&mut self, id: &str) {
145        self.tool_calls.remove(id);
146        self.tool_order.retain(|tool_id| tool_id != id);
147        self.sub_agents.remove(id);
148    }
149
150    pub fn render_tool(&self, id: &str, context: &ViewContext) -> Vec<Line> {
151        render_tool_tree(id, &self.tool_calls, &self.sub_agents, self.tick, context)
152    }
153
154    /// Clear all tracked tool calls (e.g., after pushing to scrollback).
155    pub fn clear(&mut self) {
156        self.tool_order.clear();
157        self.tool_calls.clear();
158        self.sub_agents.clear();
159    }
160
161    fn top_level_counts(&self) -> (usize, usize) {
162        let total = self
163            .tool_order
164            .iter()
165            .filter(|id| !self.sub_agents.has_sub_agents(id))
166            .count();
167        let completed = self
168            .tool_order
169            .iter()
170            .filter(|id| !self.sub_agents.has_sub_agents(id))
171            .filter_map(|id| self.tool_calls.get(id))
172            .filter(|tc| !matches!(tc.status, ToolCallStatus::Running))
173            .count();
174        (completed, total)
175    }
176
177    fn any_running_including_subagents(&self) -> bool {
178        self.tool_calls
179            .values()
180            .any(|tc| matches!(tc.status, ToolCallStatus::Running))
181            || self.sub_agents.any_running()
182    }
183}
184
185impl Default for ToolCallStatuses {
186    fn default() -> Self {
187        Self::new()
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194    use acp_utils::notifications::{SubAgentEvent, SubAgentProgressParams};
195    use tui::{DiffLine, DiffPreview, DiffTag, SplitDiffCell, SplitDiffRow};
196
197    fn ctx() -> ViewContext {
198        ViewContext::new((80, 24))
199    }
200
201    fn make_tool_call(id: &str, title: &str, raw_input: Option<&str>) -> acp::ToolCall {
202        let mut tc = acp::ToolCall::new(id.to_string(), title);
203        if let Some(input) = raw_input {
204            tc = tc.raw_input(serde_json::from_str::<serde_json::Value>(input).unwrap());
205        }
206        tc
207    }
208
209    fn make_tool_call_update(id: &str, status: acp::ToolCallStatus) -> acp::ToolCallUpdate {
210        acp::ToolCallUpdate::new(
211            id.to_string(),
212            acp::ToolCallUpdateFields::new().status(status),
213        )
214    }
215
216    fn make_sub_agent_notification(
217        parent_tool_id: &str,
218        agent_name: &str,
219        event_json: &str,
220    ) -> SubAgentProgressParams {
221        make_sub_agent_notification_with_task_id(parent_tool_id, agent_name, agent_name, event_json)
222    }
223
224    fn make_sub_agent_notification_with_task_id(
225        parent_tool_id: &str,
226        task_id: &str,
227        agent_name: &str,
228        event_json: &str,
229    ) -> SubAgentProgressParams {
230        let json = format!(
231            r#"{{"parent_tool_id":"{parent_tool_id}","task_id":"{task_id}","agent_name":"{agent_name}","event":{event_json}}}"#,
232        );
233        serde_json::from_str(&json).unwrap()
234    }
235
236    #[test]
237    fn progress_reports_sub_agent_running_tools() {
238        let mut statuses = ToolCallStatuses::new();
239        statuses.on_tool_call(&make_tool_call("parent-1", "spawn_subagent", None));
240        statuses.on_tool_call_update(&make_tool_call_update(
241            "parent-1",
242            acp::ToolCallStatus::Completed,
243        ));
244        statuses.on_sub_agent_progress(&make_sub_agent_notification(
245            "parent-1",
246            "explorer",
247            r#"{"ToolCall":{"request":{"id":"c1","name":"grep","arguments":"{}"},"model_name":"m"}}"#,
248        ));
249
250        assert!(statuses.progress().running_any);
251    }
252
253    #[test]
254    fn remove_tool_cleans_up_sub_agent_state() {
255        let mut statuses = ToolCallStatuses::new();
256        statuses.on_tool_call(&make_tool_call("parent-1", "spawn_subagent", None));
257        statuses.on_sub_agent_progress(&make_sub_agent_notification(
258            "parent-1",
259            "explorer",
260            r#"{"ToolCall":{"request":{"id":"c1","name":"grep","arguments":"{}"},"model_name":"m"}}"#,
261        ));
262
263        statuses.remove_tool("parent-1");
264        assert!(!statuses.progress().running_any);
265        assert!(statuses.render_tool("parent-1", &ctx()).is_empty());
266    }
267
268    #[test]
269    fn clear_removes_sub_agent_state() {
270        let mut statuses = ToolCallStatuses::new();
271        statuses.on_tool_call(&make_tool_call("parent-1", "spawn_subagent", None));
272        statuses.on_sub_agent_progress(&make_sub_agent_notification(
273            "parent-1",
274            "explorer",
275            r#"{"ToolCall":{"request":{"id":"c1","name":"grep","arguments":"{}"},"model_name":"m"}}"#,
276        ));
277
278        statuses.clear();
279        assert!(!statuses.progress().running_any);
280    }
281
282    #[test]
283    fn deserialize_tool_call_event() {
284        let n = make_sub_agent_notification(
285            "p1",
286            "explorer",
287            r#"{"ToolCall":{"request":{"id":"c1","name":"grep","arguments":"{\"pattern\":\"test\"}"},"model_name":"m"}}"#,
288        );
289        assert!(matches!(n.event, SubAgentEvent::ToolCall { .. }));
290    }
291
292    #[test]
293    fn deserialize_tool_call_update_event() {
294        let n = make_sub_agent_notification(
295            "p1",
296            "explorer",
297            r#"{"ToolCallUpdate":{"update":{"id":"c1","chunk":"{\"pattern\":\"updated\"}"},"model_name":"m"}}"#,
298        );
299        assert!(matches!(n.event, SubAgentEvent::ToolCallUpdate { .. }));
300    }
301
302    #[test]
303    fn deserialize_tool_result_event() {
304        let n = make_sub_agent_notification(
305            "p1",
306            "explorer",
307            r#"{"ToolResult":{"result":{"id":"c1","name":"grep","arguments":"{}","result":"ok"},"model_name":"m"}}"#,
308        );
309        assert!(matches!(n.event, SubAgentEvent::ToolResult { .. }));
310    }
311
312    #[test]
313    fn deserialize_done_event() {
314        let n = make_sub_agent_notification("p1", "explorer", r#""Done""#);
315        assert!(matches!(n.event, SubAgentEvent::Done));
316    }
317
318    #[test]
319    fn deserialize_other_variant() {
320        let n = make_sub_agent_notification("p1", "explorer", r#""Other""#);
321        assert!(matches!(n.event, SubAgentEvent::Other));
322    }
323
324    #[test]
325    fn test_diff_preview_rendered_on_success() {
326        let mut statuses = ToolCallStatuses::new();
327        statuses.on_tool_call(&make_tool_call("tool-1", "Edit", None));
328
329        let tc = statuses.tool_calls.get_mut("tool-1").unwrap();
330        tc.status = ToolCallStatus::Success;
331        tc.diff_preview = Some(DiffPreview {
332            lines: vec![
333                DiffLine {
334                    tag: DiffTag::Removed,
335                    content: "old line".to_string(),
336                },
337                DiffLine {
338                    tag: DiffTag::Added,
339                    content: "new line".to_string(),
340                },
341            ],
342            rows: vec![SplitDiffRow {
343                left: Some(SplitDiffCell {
344                    tag: DiffTag::Removed,
345                    content: "old line".to_string(),
346                    line_number: Some(1),
347                }),
348                right: Some(SplitDiffCell {
349                    tag: DiffTag::Added,
350                    content: "new line".to_string(),
351                    line_number: Some(1),
352                }),
353            }],
354            lang_hint: "rs".to_string(),
355            start_line: Some(1),
356        });
357
358        let lines = statuses.render_tool("tool-1", &ctx());
359        assert!(lines.len() > 1);
360        let all_text: String = lines.iter().map(|l| l.plain_text()).collect();
361        assert!(
362            all_text.contains("old line"),
363            "Expected removed line: {all_text}"
364        );
365        assert!(
366            all_text.contains("new line"),
367            "Expected added line: {all_text}"
368        );
369    }
370
371    #[test]
372    fn test_diff_preview_not_rendered_while_running() {
373        let mut statuses = ToolCallStatuses::new();
374        statuses.on_tool_call(&make_tool_call("tool-1", "Edit", None));
375
376        let tc = statuses.tool_calls.get_mut("tool-1").unwrap();
377        tc.diff_preview = Some(DiffPreview {
378            lines: vec![DiffLine {
379                tag: DiffTag::Added,
380                content: "new line".to_string(),
381            }],
382            rows: vec![SplitDiffRow {
383                left: None,
384                right: Some(SplitDiffCell {
385                    tag: DiffTag::Added,
386                    content: "new line".to_string(),
387                    line_number: Some(1),
388                }),
389            }],
390            lang_hint: "rs".to_string(),
391            start_line: Some(1),
392        });
393
394        let lines = statuses.render_tool("tool-1", &ctx());
395        assert_eq!(lines.len(), 1, "Should only have status line while running");
396    }
397
398    #[test]
399    fn finalize_running_marks_top_level_tools_terminal() {
400        let mut statuses = ToolCallStatuses::new();
401        statuses.on_tool_call(&make_tool_call("tool-1", "Read", None));
402
403        statuses.finalize_running(false);
404
405        assert!(!statuses.is_tool_running("tool-1"));
406        assert!(!statuses.progress().running_any);
407        let lines = statuses.render_tool("tool-1", &ctx());
408        assert!(lines[0].plain_text().contains('✓'));
409    }
410
411    #[test]
412    fn finalize_running_marks_sub_agent_tools_terminal() {
413        let mut statuses = ToolCallStatuses::new();
414        statuses.on_sub_agent_progress(&make_sub_agent_notification(
415            "parent-1",
416            "explorer",
417            r#"{"ToolCall":{"request":{"id":"c1","name":"grep","arguments":"{}"},"model_name":"m"}}"#,
418        ));
419
420        assert!(statuses.progress().running_any);
421
422        statuses.finalize_running(true);
423
424        assert!(!statuses.progress().running_any);
425    }
426}