fix(tui/history): finalize wait tool cells when done

This commit is contained in:
James Peter
2025-09-21 09:02:11 +10:00
parent 5c19353d6c
commit 5378e55da9
4 changed files with 185 additions and 96 deletions
+59 -45
View File
@@ -1814,14 +1814,16 @@ impl ChatWidget<'_> {
fn interrupt_running_task(&mut self) {
if self.bottom_pane.is_task_running() {
let mut has_wait_running = false;
for (_, idx) in self.tools_state.running_custom_tools.iter() {
if let Some(cell) = self.history_cells.get(*idx).and_then(|c| c
for entry in self.tools_state.running_custom_tools.values() {
if let Some(idx) = self.resolve_running_tool_index(entry) {
if let Some(cell) = self.history_cells.get(idx).and_then(|c| c
.as_any()
.downcast_ref::<history_cell::RunningToolCallCell>())
{
if cell.has_title("Waiting") {
has_wait_running = true;
break;
{
if cell.has_title("Waiting") {
has_wait_running = true;
break;
}
}
}
}
@@ -2933,6 +2935,20 @@ impl ChatWidget<'_> {
}
}
fn resolve_running_tool_index(&self, entry: &RunningToolEntry) -> Option<usize> {
if let Some(pos) = self
.cell_order_seq
.iter()
.position(|key| *key == entry.order_key)
{
return Some(pos);
}
if entry.fallback_index < self.history_cells.len() {
return Some(entry.fallback_index);
}
None
}
fn history_remove_at(&mut self, idx: usize) {
if idx < self.history_cells.len() {
self.history_cells.remove(idx);
@@ -4430,7 +4446,7 @@ impl ChatWidget<'_> {
if idx < self.history_cells.len() {
self.tools_state
.running_custom_tools
.insert(ToolCallId(call_id.clone()), idx);
.insert(ToolCallId(call_id.clone()), RunningToolEntry::new(ok, idx));
}
// Update border status based on tool
@@ -4474,13 +4490,16 @@ impl ChatWidget<'_> {
Ok(content) => (true, content),
Err(error) => (false, error),
};
let entry_idx = self
let running_entry = self
.tools_state
.running_custom_tools
.remove(&ToolCallId(call_id.clone()));
let resolved_idx = running_entry
.as_ref()
.and_then(|entry| self.resolve_running_tool_index(entry));
if tool_name == "apply_patch" && success {
if let Some(idx) = entry_idx {
if let Some(idx) = resolved_idx {
if idx < self.history_cells.len() {
let is_running_tool = self.history_cells[idx]
.as_any()
@@ -4499,22 +4518,12 @@ impl ChatWidget<'_> {
if tool_name == "wait" && success {
let target = wait_target_from_params(params_string.as_ref(), &call_id);
let message = format!("Waited for {}", target);
if let Some(idx) = entry_idx {
if idx < self.history_cells.len() {
self.history_replace_at(
idx,
Box::new(history_cell::new_background_event(message.clone())),
);
} else {
let _ = self.history_insert_with_key_global(
Box::new(history_cell::new_background_event(message.clone())),
ok,
);
}
let wait_cell = history_cell::new_completed_wait_tool_call(target, duration);
if let Some(idx) = resolved_idx {
self.history_replace_at(idx, Box::new(wait_cell));
} else {
let _ = self.history_insert_with_key_global(
Box::new(history_cell::new_background_event(message.clone())),
Box::new(wait_cell),
ok,
);
}
@@ -4535,15 +4544,8 @@ impl ChatWidget<'_> {
HistoryCellType::Error,
);
if let Some(idx) = entry_idx {
if idx < self.history_cells.len() {
self.history_replace_at(idx, Box::new(wait_cancelled_cell));
} else {
let _ = self.history_insert_with_key_global(
Box::new(wait_cancelled_cell),
ok,
);
}
if let Some(idx) = resolved_idx {
self.history_replace_at(idx, Box::new(wait_cancelled_cell));
} else {
let _ = self.history_insert_with_key_global(
Box::new(wait_cancelled_cell),
@@ -4565,12 +4567,8 @@ impl ChatWidget<'_> {
success,
content,
);
if let Some(idx) = entry_idx {
if idx < self.history_cells.len() {
self.history_replace_at(idx, Box::new(completed));
} else {
let _ = self.history_insert_with_key_global(Box::new(completed), ok);
}
if let Some(idx) = resolved_idx {
self.history_replace_at(idx, Box::new(completed));
} else {
let _ = self.history_insert_with_key_global(Box::new(completed), ok);
}
@@ -4588,12 +4586,8 @@ impl ChatWidget<'_> {
success,
content,
);
if let Some(idx) = entry_idx {
if idx < self.history_cells.len() {
self.history_replace_at(idx, Box::new(completed));
} else {
let _ = self.history_insert_with_key_global(Box::new(completed), ok);
}
if let Some(idx) = resolved_idx {
self.history_replace_at(idx, Box::new(completed));
} else {
let _ = self.history_insert_with_key_global(Box::new(completed), ok);
}
@@ -7106,6 +7100,11 @@ impl ChatWidget<'_> {
.downcast_mut::<history_cell::StreamingContentCell>()
{
stream.retint(&old, &new);
} else if let Some(wait) = cell
.as_any_mut()
.downcast_mut::<history_cell::WaitStatusCell>()
{
wait.retint(&old, &new);
} else if let Some(assist) = cell
.as_any_mut()
.downcast_mut::<history_cell::AssistantMarkdownCell>()
@@ -14065,9 +14064,24 @@ impl ExecState {
}
}
#[derive(Clone, Copy, Debug)]
pub(super) struct RunningToolEntry {
order_key: OrderKey,
fallback_index: usize,
}
impl RunningToolEntry {
fn new(order_key: OrderKey, fallback_index: usize) -> Self {
Self {
order_key,
fallback_index,
}
}
}
#[derive(Default)]
struct ToolState {
running_custom_tools: HashMap<ToolCallId, usize>,
running_custom_tools: HashMap<ToolCallId, RunningToolEntry>,
running_web_search: HashMap<ToolCallId, (usize, Option<String>)>,
}
#[derive(Default)]
+45 -41
View File
@@ -109,41 +109,43 @@ pub(super) fn finalize_all_running_as_interrupted(chat: &mut ChatWidget<'_>) {
}
if !chat.tools_state.running_custom_tools.is_empty() {
let entries: Vec<(super::ToolCallId, usize)> = chat
let entries: Vec<(super::ToolCallId, super::RunningToolEntry)> = chat
.tools_state
.running_custom_tools
.iter()
.map(|(k, i)| (k.clone(), *i))
.map(|(k, entry)| (k.clone(), *entry))
.collect();
for (_k, idx) in entries {
if idx < chat.history_cells.len() {
let wait_cancel_cell = Box::new(history_cell::PlainHistoryCell {
lines: vec![Line::styled(
"Wait cancelled",
Style::default()
.fg(crate::colors::error())
.add_modifier(Modifier::BOLD),
)],
kind: history_cell::HistoryCellType::Error,
});
for (_k, entry) in entries {
if let Some(idx) = chat.resolve_running_tool_index(&entry) {
if idx < chat.history_cells.len() {
let wait_cancel_cell = Box::new(history_cell::PlainHistoryCell::new(
vec![Line::styled(
"Wait cancelled",
Style::default()
.fg(crate::colors::error())
.add_modifier(Modifier::BOLD),
)],
history_cell::HistoryCellType::Error,
));
let replaced = chat.history_cells[idx]
.as_any()
.downcast_ref::<history_cell::RunningToolCallCell>()
.map(|cell| cell.has_title("Waiting"))
.unwrap_or(false);
let replaced = chat.history_cells[idx]
.as_any()
.downcast_ref::<history_cell::RunningToolCallCell>()
.map(|cell| cell.has_title("Waiting"))
.unwrap_or(false);
if replaced {
chat.history_replace_at(idx, wait_cancel_cell);
} else {
let completed = history_cell::new_completed_custom_tool_call(
"custom".to_string(),
None,
std::time::Duration::from_millis(0),
false,
"Cancelled by user.".to_string(),
);
chat.history_replace_at(idx, Box::new(completed));
if replaced {
chat.history_replace_at(idx, wait_cancel_cell);
} else {
let completed = history_cell::new_completed_custom_tool_call(
"custom".to_string(),
None,
std::time::Duration::from_millis(0),
false,
"Cancelled by user.".to_string(),
);
chat.history_replace_at(idx, Box::new(completed));
}
}
}
}
@@ -245,22 +247,24 @@ pub(super) fn finalize_all_running_due_to_answer(chat: &mut ChatWidget<'_>) {
}
if !chat.tools_state.running_custom_tools.is_empty() {
let entries: Vec<(super::ToolCallId, usize)> = chat
let entries: Vec<(super::ToolCallId, super::RunningToolEntry)> = chat
.tools_state
.running_custom_tools
.iter()
.map(|(k, i)| (k.clone(), *i))
.map(|(k, entry)| (k.clone(), *entry))
.collect();
for (_k, idx) in entries {
if idx < chat.history_cells.len() {
let completed = history_cell::new_completed_custom_tool_call(
"custom".to_string(),
None,
std::time::Duration::from_millis(0),
true,
"Final answer received".to_string(),
);
chat.history_replace_at(idx, Box::new(completed));
for (_k, entry) in entries {
if let Some(idx) = chat.resolve_running_tool_index(&entry) {
if idx < chat.history_cells.len() {
let completed = history_cell::new_completed_custom_tool_call(
"custom".to_string(),
None,
std::time::Duration::from_millis(0),
true,
"Final answer received".to_string(),
);
chat.history_replace_at(idx, Box::new(completed));
}
}
}
chat.tools_state.running_custom_tools.clear();
+7 -5
View File
@@ -66,17 +66,19 @@ pub(super) fn mcp_begin(chat: &mut ChatWidget<'_>, ev: McpToolCallBeginEvent, ke
let idx = chat.history_insert_with_key_global(Box::new(cell), key);
chat.tools_state
.running_custom_tools
.insert(super::ToolCallId(call_id), idx);
.insert(super::ToolCallId(call_id), super::RunningToolEntry::new(key, idx));
}
pub(super) fn mcp_end(chat: &mut ChatWidget<'_>, ev: McpToolCallEndEvent, key: OrderKey) {
let McpToolCallEndEvent { call_id, duration, invocation, result } = ev;
let success = !result.as_ref().map(|r| r.is_error.unwrap_or(false)).unwrap_or(false);
let completed = history_cell::new_completed_mcp_tool_call(80, invocation, duration, success, result);
if let Some(idx) = chat.tools_state.running_custom_tools.remove(&super::ToolCallId(call_id)) {
if idx < chat.history_cells.len() {
chat.history_replace_at(idx, completed);
return;
if let Some(entry) = chat.tools_state.running_custom_tools.remove(&super::ToolCallId(call_id)) {
if let Some(idx) = chat.resolve_running_tool_index(&entry) {
if idx < chat.history_cells.len() {
chat.history_replace_at(idx, completed);
return;
}
}
}
let _ = chat.history_insert_with_key_global(Box::new(completed), key);
+74 -5
View File
@@ -4043,6 +4043,69 @@ impl HistoryCell for RunningToolCallCell {
}
}
// ==================== WaitStatusCell (completed wait) ====================
pub(crate) struct WaitStatusCell {
lines: Vec<Line<'static>>,
}
impl WaitStatusCell {
pub(crate) fn retint(&mut self, old: &crate::theme::Theme, new: &crate::theme::Theme) {
retint_lines_in_place(&mut self.lines, old, new);
}
}
impl HistoryCell for WaitStatusCell {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
self
}
fn kind(&self) -> HistoryCellType {
HistoryCellType::Plain
}
fn display_lines(&self) -> Vec<Line<'static>> {
self.lines.clone()
}
fn gutter_symbol(&self) -> Option<&'static str> {
Some("◓")
}
}
pub(crate) fn new_completed_wait_tool_call(target: String, duration: Duration) -> WaitStatusCell {
let mut lines: Vec<Line<'static>> = Vec::new();
let mut duration_str = format_duration(duration);
if duration_str.ends_with(" 00s") {
duration_str.truncate(duration_str.len() - 4);
}
let mut header_spans = Vec::new();
header_spans.push(
Span::styled(
"Waited",
Style::default()
.fg(crate::colors::success())
.add_modifier(Modifier::BOLD),
),
);
header_spans.push(Span::styled(
format!(" ({duration_str})"),
Style::default().fg(crate::colors::text_dim()),
));
lines.push(Line::from(header_spans));
if !target.is_empty() {
lines.push(Line::styled(
format!("for {target}"),
Style::default().fg(crate::colors::text_dim()),
));
}
lines.push(Line::from(""));
WaitStatusCell { lines }
}
impl RunningToolCallCell {
fn strip_zero_seconds_suffix(mut duration: String) -> String {
if duration.ends_with(" 00s") {
@@ -7893,11 +7956,17 @@ pub(crate) fn new_running_custom_tool_call(
));
arg_lines.push(Line::from(spans));
wait_has_target = true;
} else if let Some(cid) = json.get("call_id").and_then(|v| v.as_str()) {
arg_lines.push(Line::from(vec![
Span::styled("└ call_id: ", Style::default().fg(crate::colors::text_dim())),
Span::styled(cid.to_string(), Style::default().fg(crate::colors::text())),
]));
}
if let Some(cid) = json.get("call_id").and_then(|v| v.as_str()) {
if !wait_has_target {
arg_lines.push(Line::from(vec![
Span::styled(
"└ call_id: ",
Style::default().fg(crate::colors::text_dim()),
),
Span::styled(cid.to_string(), Style::default().fg(crate::colors::text())),
]));
}
wait_has_call_id = true;
}
} else {