Skip to main content

wisp/conversation/
tool_calls.rs

1use crate::git_review::FileDiff;
2use acp_utils::AETHER_TOOL_NAME_META_KEY;
3use acp_utils::notifications::{SubAgentEvent, SubAgentProgressParams};
4use agent_client_protocol::schema::v1 as acp;
5
6pub const SUB_AGENT_VISIBLE_TOOL_LIMIT: usize = 3;
7
8/// A tracked tool call within a sub-agent.
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct SubAgentToolCall {
11    pub id: String,
12    pub name: String,
13    pub arguments: String,
14    pub raw_input: String,
15    pub display_value: Option<String>,
16    pub status: ToolStatus,
17    kind: ToolKind,
18}
19
20impl SubAgentToolCall {
21    pub fn bash_command(&self) -> Option<String> {
22        bash_command(self.kind, &self.raw_input)
23    }
24}
25
26/// Per-sub-agent state: tracks its tool calls in arrival order.
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct SubAgentState {
29    pub task_id: String,
30    pub agent_name: String,
31    pub done: bool,
32    pub tool_calls: Vec<SubAgentToolCall>,
33}
34
35impl SubAgentState {
36    fn tool_call_mut(&mut self, id: &str) -> Option<&mut SubAgentToolCall> {
37        self.tool_calls.iter_mut().find(|call| call.id == id)
38    }
39
40    /// The call with `id`, appending a running placeholder when it is the first
41    /// event seen for it.
42    fn upsert(&mut self, id: &str, name: &str, arguments: String) -> &mut SubAgentToolCall {
43        let index = self.tool_calls.iter().position(|call| call.id == id).unwrap_or_else(|| {
44            self.tool_calls.push(SubAgentToolCall {
45                id: id.to_string(),
46                name: name.to_string(),
47                raw_input: arguments.clone(),
48                arguments,
49                display_value: None,
50                status: ToolStatus::Running,
51                kind: tool_kind(name),
52            });
53            self.tool_calls.len() - 1
54        });
55        &mut self.tool_calls[index]
56    }
57}
58
59fn apply_sub_agent_progress(states: &mut Vec<SubAgentState>, notification: &SubAgentProgressParams) {
60    let index = states.iter().position(|agent| agent.task_id == notification.task_id).unwrap_or_else(|| {
61        states.push(SubAgentState {
62            task_id: notification.task_id.clone(),
63            agent_name: notification.agent_name.clone(),
64            done: false,
65            tool_calls: Vec::new(),
66        });
67        states.len() - 1
68    });
69    let agent = &mut states[index];
70
71    match &notification.event {
72        SubAgentEvent::ToolCall { request } => {
73            let call = agent.upsert(&request.id, &request.name, request.arguments.clone());
74            update_title(&mut call.name, &request.name);
75            call.kind = tool_kind(&request.name);
76            call.arguments.clone_from(&request.arguments);
77            call.raw_input.clone_from(&request.arguments);
78            call.status = ToolStatus::Running;
79        }
80        SubAgentEvent::ToolCallUpdate { update } => {
81            let call = agent.upsert(&update.id, "tool", String::new());
82            call.arguments.push_str(&update.chunk);
83            call.raw_input.push_str(&update.chunk);
84            call.status = ToolStatus::Running;
85        }
86        SubAgentEvent::ToolResult { result } => {
87            if let Some(call) = agent.tool_call_mut(&result.id) {
88                call.status = ToolStatus::Success;
89                if let Some(result_meta) = &result.result_meta {
90                    call.name.clone_from(&result_meta.display.title);
91                    call.display_value = Some(result_meta.display.value.clone());
92                }
93            }
94        }
95        SubAgentEvent::ToolError { error } => {
96            if let Some(call) = agent.tool_call_mut(&error.id) {
97                call.status = ToolStatus::Error("failed".to_string());
98            }
99        }
100        SubAgentEvent::Done => agent.done = true,
101        SubAgentEvent::Other => {}
102    }
103}
104
105#[derive(Debug, Clone, PartialEq, Eq)]
106pub struct ToolCall {
107    pub id: String,
108    pub title: String,
109    pub status: ToolStatus,
110    pub diff: Option<Box<FileDiff>>,
111    pub raw_input: String,
112    pub display_value: Option<String>,
113    pub sub_agents: Vec<SubAgentState>,
114    kind: ToolKind,
115}
116
117impl ToolCall {
118    pub(crate) fn from_acp(tool_call: &acp::ToolCall) -> Self {
119        let tool_name = tool_call
120            .meta
121            .as_ref()
122            .and_then(|meta| meta.get(AETHER_TOOL_NAME_META_KEY))
123            .and_then(serde_json::Value::as_str)
124            .unwrap_or(&tool_call.title);
125        let raw_input = tool_call.raw_input.as_ref().map_or_else(String::new, raw_input_fragment);
126        Self {
127            id: tool_call.tool_call_id.0.to_string(),
128            title: tool_call.title.clone(),
129            status: ToolStatus::Running,
130            diff: None,
131            raw_input,
132            display_value: None,
133            sub_agents: Vec::new(),
134            kind: tool_kind(tool_name),
135        }
136    }
137
138    pub(crate) fn apply_update(&mut self, update: &acp::ToolCallUpdate) {
139        if let Some(title) = &update.fields.title {
140            update_title(&mut self.title, title);
141        }
142        if let Some(raw_input) = &update.fields.raw_input {
143            self.raw_input.push_str(&raw_input_fragment(raw_input));
144        }
145        if let Some(meta) = &update.meta
146            && let Some(serde_json::Value::String(display_value)) = meta.get("display_value")
147        {
148            self.display_value = Some(display_value.clone());
149        }
150        if let Some(content) = &update.fields.content {
151            for item in content {
152                if let acp::ToolCallContent::Diff(diff) = item {
153                    self.diff = Some(Box::new(FileDiff::from_texts(
154                        diff.path.display().to_string(),
155                        diff.old_text.as_deref().unwrap_or_default(),
156                        &diff.new_text,
157                    )));
158                }
159            }
160        }
161        if let Some(status) = update.fields.status {
162            match status {
163                acp::ToolCallStatus::Completed => self.status = ToolStatus::Success,
164                acp::ToolCallStatus::Failed => self.status = ToolStatus::Error("failed".to_string()),
165                acp::ToolCallStatus::InProgress | acp::ToolCallStatus::Pending => self.status = ToolStatus::Running,
166                _ => {}
167            }
168        }
169    }
170
171    pub(crate) fn apply_sub_agent_progress(&mut self, notification: &SubAgentProgressParams) {
172        apply_sub_agent_progress(&mut self.sub_agents, notification);
173    }
174
175    pub(crate) fn finalize(&mut self, terminal_status: &ToolStatus) {
176        if self.status == ToolStatus::Running {
177            self.status = terminal_status.clone();
178        }
179        for agent in &mut self.sub_agents {
180            agent.done = true;
181            for call in &mut agent.tool_calls {
182                if matches!(call.status, ToolStatus::Running) {
183                    call.status = terminal_status.clone();
184                }
185            }
186        }
187    }
188
189    pub fn bash_command(&self) -> Option<String> {
190        bash_command(self.kind, &self.raw_input)
191    }
192
193    pub(crate) fn is_running(&self) -> bool {
194        self.status == ToolStatus::Running
195            || self.sub_agents.iter().any(|agent| {
196                !agent.done || agent.tool_calls.iter().any(|call| matches!(call.status, ToolStatus::Running))
197            })
198    }
199
200    /// Whether this call's rendering can no longer change: it reached a
201    /// terminal status and every spawned sub-agent has finished. A background
202    /// spawn completes before its agents start reporting, so an empty tree on
203    /// a completed spawner means "not yet", not "none".
204    pub(crate) fn rendering_final(&self) -> bool {
205        !self.is_running() && (self.kind != ToolKind::SpawnSubagent || !self.sub_agents.is_empty())
206    }
207}
208
209#[derive(Debug, Clone, PartialEq, Eq)]
210pub enum ToolStatus {
211    Running,
212    Success,
213    Error(String),
214}
215
216fn update_title(current: &mut String, new_title: &str) {
217    if !new_title.is_empty() {
218        current.clear();
219        current.push_str(new_title);
220    }
221}
222
223pub(crate) fn raw_input_fragment(raw_input: &serde_json::Value) -> String {
224    raw_input.as_str().map_or_else(|| raw_input.to_string(), str::to_string)
225}
226
227#[derive(Debug, Clone, Copy, PartialEq, Eq)]
228enum ToolKind {
229    Bash,
230    SpawnSubagent,
231    Other,
232}
233
234fn tool_kind(tool_name: &str) -> ToolKind {
235    let name = tool_name.rsplit("__").next().unwrap_or(tool_name);
236    if name.eq_ignore_ascii_case("bash") {
237        ToolKind::Bash
238    } else if name.eq_ignore_ascii_case("spawn_subagent") {
239        ToolKind::SpawnSubagent
240    } else {
241        ToolKind::Other
242    }
243}
244
245fn bash_command(kind: ToolKind, raw_input: &str) -> Option<String> {
246    if kind != ToolKind::Bash {
247        return None;
248    }
249    serde_json::from_str::<serde_json::Value>(raw_input).ok()?.get("command")?.as_str().map(str::to_string)
250}