Skip to main content

wisp/conversation/
tool_calls.rs

1use acp_utils::notifications::{SubAgentEvent, SubAgentProgressParams};
2use agent_client_protocol::schema::{MaybeUndefined, v2 as acp};
3
4pub const SUB_AGENT_VISIBLE_TOOL_LIMIT: usize = 3;
5
6/// A tracked tool call within a sub-agent.
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct SubAgentToolCall {
9    pub id: String,
10    pub name: String,
11    pub raw_input: String,
12    pub display_value: Option<String>,
13    pub status: ToolStatus,
14    kind: ToolKind,
15}
16
17impl SubAgentToolCall {
18    pub fn bash_command(&self) -> Option<String> {
19        bash_command(self.kind, &self.raw_input)
20    }
21}
22
23/// Per-sub-agent state: tracks its tool calls in arrival order.
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct SubAgentState {
26    pub task_id: String,
27    pub agent_name: String,
28    pub done: bool,
29    pub tool_calls: Vec<SubAgentToolCall>,
30}
31
32impl SubAgentState {
33    fn tool_call_mut(&mut self, id: &str) -> Option<&mut SubAgentToolCall> {
34        self.tool_calls.iter_mut().find(|call| call.id == id)
35    }
36
37    /// The call with `id`, appending a running placeholder when it is the first
38    /// event seen for it.
39    fn upsert(&mut self, id: &str, name: &str, arguments: String) -> &mut SubAgentToolCall {
40        let index = self.tool_calls.iter().position(|call| call.id == id).unwrap_or_else(|| {
41            self.tool_calls.push(SubAgentToolCall {
42                id: id.to_string(),
43                name: name.to_string(),
44                raw_input: arguments,
45                display_value: None,
46                status: ToolStatus::Running,
47                kind: tool_kind(name),
48            });
49            self.tool_calls.len() - 1
50        });
51        &mut self.tool_calls[index]
52    }
53}
54
55#[derive(Debug, Clone, PartialEq)]
56pub struct ToolCall {
57    pub status: ToolStatus,
58    pub sub_agents: Vec<SubAgentState>,
59    protocol: Box<acp::ToolCallUpdate>,
60}
61
62impl ToolCall {
63    pub fn from_update(update: &acp::ToolCallUpdate) -> Self {
64        let mut tool = Self {
65            status: ToolStatus::Running,
66            sub_agents: Vec::new(),
67            protocol: Box::new(update.clone()),
68        };
69        tool.refresh_status();
70        tool
71    }
72
73    pub fn title(&self) -> &str {
74        self.protocol.title.value().map_or("", String::as_str)
75    }
76
77    pub fn raw_input(&self) -> String {
78        self.protocol.raw_input.value().map_or_else(String::new, raw_input_fragment)
79    }
80
81    pub fn display_value(&self) -> Option<&str> {
82        self.meta_str("display_value")
83    }
84
85    pub fn content(&self) -> &[acp::ToolCallContent] {
86        self.protocol.content.value().map_or(&[], Vec::as_slice)
87    }
88
89    pub fn diffs(&self) -> impl Iterator<Item = &acp::Diff> {
90        self.content().iter().filter_map(|content| match content {
91            acp::ToolCallContent::Diff(diff) => Some(diff),
92            _ => None,
93        })
94    }
95
96    pub fn apply_update(&mut self, update: &acp::ToolCallUpdate) {
97        self.protocol.apply_update(update.clone());
98        self.refresh_status();
99    }
100
101    pub fn append_content(&mut self, content: acp::ToolCallContent) {
102        match &mut self.protocol.content {
103            MaybeUndefined::Value(items) => items.push(content),
104            value => *value = MaybeUndefined::Value(vec![content]),
105        }
106    }
107
108    pub(crate) fn apply_sub_agent_progress(&mut self, notification: &SubAgentProgressParams) {
109        apply_sub_agent_progress(&mut self.sub_agents, notification);
110    }
111
112    pub(crate) fn finalize(&mut self, terminal_status: &ToolStatus) {
113        if self.status == ToolStatus::Running {
114            self.status = terminal_status.clone();
115        }
116        for agent in &mut self.sub_agents {
117            agent.done = true;
118            for call in &mut agent.tool_calls {
119                if matches!(call.status, ToolStatus::Running) {
120                    call.status = terminal_status.clone();
121                }
122            }
123        }
124    }
125
126    pub fn bash_command(&self) -> Option<String> {
127        bash_command(self.kind(), &self.raw_input())
128    }
129
130    pub(crate) fn is_running(&self) -> bool {
131        self.status == ToolStatus::Running
132            || self.sub_agents.iter().any(|agent| {
133                !agent.done || agent.tool_calls.iter().any(|call| matches!(call.status, ToolStatus::Running))
134            })
135    }
136
137    /// Whether this call can enter native history: it reached a terminal
138    /// status and every spawned sub-agent has finished. A background
139    /// spawn completes before its agents start reporting, so an empty tree on
140    /// a completed spawner means "not yet", not "none".
141    pub(crate) fn rendering_final(&self) -> bool {
142        !self.is_running() && (self.kind() != ToolKind::SpawnSubagent || !self.sub_agents.is_empty())
143    }
144
145    fn kind(&self) -> ToolKind {
146        tool_kind(self.protocol.name.value().map_or_else(|| self.title(), String::as_str))
147    }
148
149    /// Re-derives the coarse status from the merged protocol update; `Undefined`
150    /// fields keep their previous value, so re-running this is idempotent.
151    fn refresh_status(&mut self) {
152        self.status = match self.protocol.status.value() {
153            Some(acp::ToolCallStatus::Completed) => ToolStatus::Success,
154            Some(acp::ToolCallStatus::Failed) => ToolStatus::Error("failed".to_string()),
155            Some(acp::ToolCallStatus::Cancelled) => ToolStatus::Error("cancelled".to_string()),
156            _ => ToolStatus::Running,
157        };
158    }
159
160    fn meta_str(&self, key: &str) -> Option<&str> {
161        self.protocol.meta.value().and_then(|meta| meta.get(key)).and_then(serde_json::Value::as_str)
162    }
163}
164
165#[derive(Debug, Clone, PartialEq, Eq)]
166pub enum ToolStatus {
167    Running,
168    Success,
169    Error(String),
170}
171
172fn apply_sub_agent_progress(states: &mut Vec<SubAgentState>, notification: &SubAgentProgressParams) {
173    let index = states.iter().position(|agent| agent.task_id == notification.task_id).unwrap_or_else(|| {
174        states.push(SubAgentState {
175            task_id: notification.task_id.clone(),
176            agent_name: notification.agent_name.clone(),
177            done: false,
178            tool_calls: Vec::new(),
179        });
180        states.len() - 1
181    });
182    let agent = &mut states[index];
183
184    match &notification.event {
185        SubAgentEvent::ToolCall { request } => {
186            let call = agent.upsert(&request.id, &request.name, request.arguments.clone());
187            update_title(&mut call.name, &request.name);
188            call.kind = tool_kind(&request.name);
189            call.raw_input.clone_from(&request.arguments);
190            call.status = ToolStatus::Running;
191        }
192        SubAgentEvent::ToolCallUpdate { update } => {
193            let call = agent.upsert(&update.id, "tool", String::new());
194            call.raw_input.push_str(&update.chunk);
195            call.status = ToolStatus::Running;
196        }
197        SubAgentEvent::ToolResult { result } => {
198            if let Some(call) = agent.tool_call_mut(&result.id) {
199                call.status = ToolStatus::Success;
200                if let Some(result_meta) = &result.result_meta {
201                    call.name.clone_from(&result_meta.display.title);
202                    call.display_value = Some(result_meta.display.value.clone());
203                }
204            }
205        }
206        SubAgentEvent::ToolError { error } => {
207            if let Some(call) = agent.tool_call_mut(&error.id) {
208                call.status = ToolStatus::Error("failed".to_string());
209            }
210        }
211        SubAgentEvent::Done => agent.done = true,
212        SubAgentEvent::Other => {}
213    }
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
223fn 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}