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#[derive(Clone)]
17pub struct ToolCallStatuses {
18 tool_order: Vec<String>,
20 tool_calls: HashMap<String, TrackedToolCall>,
22 sub_agents: SubAgentTracker,
24 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 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 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 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 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 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}