improve MCP tool call styling

Author: nornagon-openai

codex-rs/tui/src/chatwidget.rs

@@ -66,6 +66,7 @@ use crate::history_cell;
 use crate::history_cell::CommandOutput;
 use crate::history_cell::ExecCell;
 use crate::history_cell::HistoryCell;
+use crate::history_cell::McpToolCallCell;
 use crate::history_cell::PatchEventType;
 use crate::slash_command::SlashCommand;
 use crate::text_formatting::truncate_text;
@@ -114,7 +115,7 @@ pub(crate) struct ChatWidget {
     app_event_tx: AppEventSender,
     codex_op_tx: UnboundedSender<Op>,
     bottom_pane: BottomPane,
-    active_exec_cell: Option<ExecCell>,
+    active_cell: Option<Box<dyn HistoryCell>>,
     config: Config,
     auth_manager: Arc<AuthManager>,
     session_header: SessionHeader,
@@ -281,7 +282,7 @@ impl ChatWidget {
     /// and stop/clear running UI state.
     fn finalize_turn_with_error_message(&mut self, message: String) {
         // Ensure any spinner is replaced by a red ✗ and flushed into history.
-        self.finalize_active_exec_cell_as_failed();
+        self.finalize_active_cell_as_failed();
         // Emit the provided error message/history cell.
         self.add_to_history(history_cell::new_error_event(message));
         // Reset running state and clear streaming buffers.
@@ -481,8 +482,8 @@ impl ChatWidget {
 
     #[inline]
     fn handle_streaming_delta(&mut self, delta: String) {
-        // Before streaming agent content, flush any active exec cell group.
-        self.flush_active_exec_cell();
+        // Before streaming agent content, flush any active cell group.
+        self.flush_active_cell();
         let sink = AppEventHistorySink(self.app_event_tx.clone());
         self.stream.begin(&sink);
         self.stream.push_and_maybe_commit(&delta, &sink);
@@ -496,16 +497,25 @@ impl ChatWidget {
             None => (vec![ev.call_id.clone()], Vec::new()),
         };
 
-        if self.active_exec_cell.is_none() {
-            // This should have been created by handle_exec_begin_now, but in case it wasn't,
-            // create it now.
-            self.active_exec_cell = Some(history_cell::new_active_exec_command(
+        let needs_new = self
+            .active_cell
+            .as_ref()
+            .map(|cell| cell.as_any().downcast_ref::<ExecCell>().is_none())
+            .unwrap_or(true);
+        if needs_new {
+            self.flush_active_cell();
+            self.active_cell = Some(Box::new(history_cell::new_active_exec_command(
                 ev.call_id.clone(),
                 command,
                 parsed,
-            ));
+            )));
         }
-        if let Some(cell) = self.active_exec_cell.as_mut() {
+
+        if let Some(cell) = self
+            .active_cell
+            .as_mut()
+            .and_then(|c| c.as_any_mut().downcast_mut::<ExecCell>())
+        {
             cell.complete_call(
                 &ev.call_id,
                 CommandOutput {
@@ -517,7 +527,7 @@ impl ChatWidget {
                 ev.duration,
             );
             if cell.should_flush() {
-                self.flush_active_exec_cell();
+                self.flush_active_cell();
             }
         }
     }
@@ -584,58 +594,76 @@ impl ChatWidget {
                 parsed_cmd: ev.parsed_cmd.clone(),
             },
         );
-        if let Some(exec) = &self.active_exec_cell {
-            if let Some(new_exec) = exec.with_added_call(
+        if let Some(cell) = self
+            .active_cell
+            .as_mut()
+            .and_then(|c| c.as_any_mut().downcast_mut::<ExecCell>())
+            && let Some(new_exec) = cell.with_added_call(
                 ev.call_id.clone(),
                 ev.command.clone(),
                 ev.parsed_cmd.clone(),
-            ) {
-                self.active_exec_cell = Some(new_exec);
-            } else {
-                // Make a new cell.
-                self.flush_active_exec_cell();
-                self.active_exec_cell = Some(history_cell::new_active_exec_command(
-                    ev.call_id.clone(),
-                    ev.command.clone(),
-                    ev.parsed_cmd,
-                ));
-            }
+            )
+        {
+            *cell = new_exec;
         } else {
-            self.active_exec_cell = Some(history_cell::new_active_exec_command(
+            self.flush_active_cell();
+
+            self.active_cell = Some(Box::new(history_cell::new_active_exec_command(
                 ev.call_id.clone(),
                 ev.command.clone(),
                 ev.parsed_cmd,
-            ));
+            )));
         }
 
-        // Request a redraw so the working header and command list are visible immediately.
         self.request_redraw();
     }
 
     pub(crate) fn handle_mcp_begin_now(&mut self, ev: McpToolCallBeginEvent) {
         self.flush_answer_stream_with_separator();
-        self.add_to_history(history_cell::new_active_mcp_tool_call(ev.invocation));
+        self.flush_active_cell();
+        self.active_cell = Some(Box::new(history_cell::new_active_mcp_tool_call(
+            ev.call_id,
+            ev.invocation,
+        )));
+        self.request_redraw();
     }
     pub(crate) fn handle_mcp_end_now(&mut self, ev: McpToolCallEndEvent) {
         self.flush_answer_stream_with_separator();
-        self.add_boxed_history(history_cell::new_completed_mcp_tool_call(
-            80,
-            ev.invocation,
-            ev.duration,
-            ev.result
-                .as_ref()
-                .map(|r| !r.is_error.unwrap_or(false))
-                .unwrap_or(false),
-            ev.result,
-        ));
+
+        let McpToolCallEndEvent {
+            call_id,
+            invocation,
+            duration,
+            result,
+        } = ev;
+
+        let extra_cell = match self
+            .active_cell
+            .as_mut()
+            .and_then(|cell| cell.as_any_mut().downcast_mut::<McpToolCallCell>())
+        {
+            Some(cell) if cell.call_id() == call_id => cell.complete(duration, result),
+            _ => {
+                self.flush_active_cell();
+                let mut cell = history_cell::new_active_mcp_tool_call(call_id, invocation);
+                let extra_cell = cell.complete(duration, result);
+                self.active_cell = Some(Box::new(cell));
+                extra_cell
+            }
+        };
+
+        self.flush_active_cell();
+        if let Some(extra) = extra_cell {
+            self.add_boxed_history(extra);
+        }
     }
 
     fn layout_areas(&self, area: Rect) -> [Rect; 3] {
         let bottom_min = self.bottom_pane.desired_height(area.width).min(area.height);
         let remaining = area.height.saturating_sub(bottom_min);
 
         let active_desired = self
-            .active_exec_cell
+            .active_cell
             .as_ref()
             .map_or(0, |c| c.desired_height(area.width) + 1);
         let active_height = active_desired.min(remaining);
@@ -680,7 +708,7 @@ impl ChatWidget {
                 placeholder_text: placeholder,
                 disable_paste_burst: config.disable_paste_burst,
             }),
-            active_exec_cell: None,
+            active_cell: None,
             config: config.clone(),
             auth_manager,
             session_header: SessionHeader::new(config.model.clone()),
@@ -736,7 +764,7 @@ impl ChatWidget {
                 placeholder_text: placeholder,
                 disable_paste_burst: config.disable_paste_burst,
             }),
-            active_exec_cell: None,
+            active_cell: None,
             config: config.clone(),
             auth_manager,
             session_header: SessionHeader::new(config.model.clone()),
@@ -762,7 +790,7 @@ impl ChatWidget {
     pub fn desired_height(&self, width: u16) -> u16 {
         self.bottom_pane.desired_height(width)
             + self
-                .active_exec_cell
+                .active_cell
                 .as_ref()
                 .map_or(0, |c| c.desired_height(width) + 1)
     }
@@ -974,10 +1002,9 @@ impl ChatWidget {
         }
     }
 
-    fn flush_active_exec_cell(&mut self) {
-        if let Some(active) = self.active_exec_cell.take() {
-            self.app_event_tx
-                .send(AppEvent::InsertHistoryCell(Box::new(active)));
+    fn flush_active_cell(&mut self) {
+        if let Some(active) = self.active_cell.take() {
+            self.app_event_tx.send(AppEvent::InsertHistoryCell(active));
         }
     }
 
@@ -988,7 +1015,7 @@ impl ChatWidget {
     fn add_boxed_history(&mut self, cell: Box<dyn HistoryCell>) {
         if !cell.display_lines(u16::MAX).is_empty() {
             // Only break exec grouping if the cell renders visible lines.
-            self.flush_active_exec_cell();
+            self.flush_active_cell();
         }
         self.app_event_tx.send(AppEvent::InsertHistoryCell(cell));
     }
@@ -1166,12 +1193,16 @@ impl ChatWidget {
         }
     }
 
-    /// Mark the active exec cell as failed (✗) and flush it into history.
-    fn finalize_active_exec_cell_as_failed(&mut self) {
-        if let Some(cell) = self.active_exec_cell.take() {
-            let cell = cell.into_failed();
-            // Insert finalized exec into history and keep grouping consistent.
-            self.add_to_history(cell);
+    /// Mark the active cell as failed (✗) and flush it into history.
+    fn finalize_active_cell_as_failed(&mut self) {
+        if let Some(mut cell) = self.active_cell.take() {
+            // Insert finalized cell into history and keep grouping consistent.
+            if let Some(exec) = cell.as_any_mut().downcast_mut::<ExecCell>() {
+                exec.mark_failed();
+            } else if let Some(tool) = cell.as_any_mut().downcast_mut::<McpToolCallCell>() {
+                tool.mark_failed();
+            }
+            self.add_boxed_history(cell);
         }
     }
 
@@ -1468,12 +1499,16 @@ impl WidgetRef for &ChatWidget {
         let [_, active_cell_area, bottom_pane_area] = self.layout_areas(area);
         (&self.bottom_pane).render(bottom_pane_area, buf);
         if !active_cell_area.is_empty()
-            && let Some(cell) = &self.active_exec_cell
+            && let Some(cell) = &self.active_cell
         {
-            let mut active_cell_area = active_cell_area;
-            active_cell_area.y = active_cell_area.y.saturating_add(1);
-            active_cell_area.height -= 1;
-            cell.render_ref(active_cell_area, buf);
+            let mut area = active_cell_area;
+            area.y = area.y.saturating_add(1);
+            area.height = area.height.saturating_sub(1);
+            if let Some(exec) = cell.as_any().downcast_ref::<ExecCell>() {
+                exec.render_ref(area, buf);
+            } else if let Some(tool) = cell.as_any().downcast_ref::<McpToolCallCell>() {
+                tool.render_ref(area, buf);
+            }
         }
     }
 }

codex-rs/tui/src/chatwidget/tests.rs

@@ -234,7 +234,7 @@ fn make_chatwidget_manual() -> (
         app_event_tx,
         codex_op_tx: op_tx,
         bottom_pane: bottom,
-        active_exec_cell: None,
+        active_cell: None,
         config: cfg.clone(),
         auth_manager,
         session_header: SessionHeader::new(cfg.model.clone()),
@@ -434,9 +434,9 @@ fn end_exec(chat: &mut ChatWidget, call_id: &str, stdout: &str, stderr: &str, ex
 
 fn active_blob(chat: &ChatWidget) -> String {
     let lines = chat
-        .active_exec_cell
+        .active_cell
         .as_ref()
-        .expect("active exec cell present")
+        .expect("active cell present")
         .display_lines(80);
     lines_to_single_string(&lines)
 }