From a67a67f3258fc21e147b6786a143fe3e15e6d5ba Mon Sep 17 00:00:00 2001
From: Reilly Wood <163153147+rgwood-dd@users.noreply.github.com>
Date: Tue, 3 Jun 2025 14:29:26 -0700
Subject: [PATCH] codex-rs: make tool calls prettier (#1211)

This PR overhauls how active tool calls and completed tool calls are
displayed:

1. More use of colour to indicate success/failure and distinguish
between components like tool name+arguments
2. Previously, the entire `CallToolResult` was serialized to JSON and
pretty-printed. Now, we extract each individual `CallToolResultContent`
and print those
1. The previous solution was wasting space by unnecessarily showing
details of the `CallToolResult` struct to users, without formatting the
actual tool call results nicely
2. We're now able to show users more information from tool results in
less space, with nicer formatting when tools return JSON results

### Before:

<img width="1251" alt="Screenshot 2025-06-03 at 11 24 26"
src="https://github.com/user-attachments/assets/5a58f222-219c-4c53-ace7-d887194e30cf"
/>

### After:

<img width="1265" alt="image"
src="https://github.com/user-attachments/assets/99fe54d0-9ebe-406a-855b-7aa529b91274"
/>

## Future Work

1. Integrate image tool result handling better. We should be able to
display images even if they're not the first `CallToolResultContent`
2. Users should have some way to view the full version of truncated tool
results
3. It would be nice to add some left padding for tool results, make it
more clear that they are results. This is doable, just a little fiddly
due to the way `first_visible_line` scrolling works
4. There's almost certainly a better way to format JSON than "all on 1
line with spaces to make Ratatui wrapping work". But I think that works
OK for now.
---
 codex-rs/Cargo.lock                           |   1 +
 codex-rs/mcp-types/src/lib.rs                 |   1 +
 codex-rs/tui/Cargo.toml                       |   3 +-
 .../tui/src/conversation_history_widget.rs    |   3 +-
 codex-rs/tui/src/history_cell.rs              | 121 +++++---
 codex-rs/tui/src/lib.rs                       |   1 +
 codex-rs/tui/src/text_formatting.rs           | 268 ++++++++++++++++++
 7 files changed, 350 insertions(+), 48 deletions(-)
 create mode 100644 codex-rs/tui/src/text_formatting.rs

diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock
index 694e1138..ae574ba3 100644
--- a/codex-rs/Cargo.lock
+++ b/codex-rs/Cargo.lock
@@ -768,6 +768,7 @@ dependencies = [
  "tui-input",
  "tui-markdown",
  "tui-textarea",
+ "unicode-segmentation",
  "uuid",
 ]
 
diff --git a/codex-rs/mcp-types/src/lib.rs b/codex-rs/mcp-types/src/lib.rs
index afd6f4ad..0ed51853 100644
--- a/codex-rs/mcp-types/src/lib.rs
+++ b/codex-rs/mcp-types/src/lib.rs
@@ -1144,6 +1144,7 @@ pub enum ServerRequest {
 
 #[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
 #[serde(untagged)]
+#[allow(clippy::large_enum_variant)]
 pub enum ServerResult {
     Result(Result),
     InitializeResult(InitializeResult),
diff --git a/codex-rs/tui/Cargo.toml b/codex-rs/tui/Cargo.toml
index 235f5f0c..ffc107e8 100644
--- a/codex-rs/tui/Cargo.toml
+++ b/codex-rs/tui/Cargo.toml
@@ -34,7 +34,7 @@ ratatui = { version = "0.29.0", features = [
 ] }
 ratatui-image = "8.0.0"
 regex-lite = "0.1"
-serde_json = "1"
+serde_json = { version = "1", features = ["preserve_order"] }
 shlex = "1.3.0"
 strum = "0.27.1"
 strum_macros = "0.27.1"
@@ -51,6 +51,7 @@ tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
 tui-input = "0.11.1"
 tui-markdown = "0.3.3"
 tui-textarea = "0.7.0"
+unicode-segmentation = "1.12.0"
 uuid = "1"
 
 [dev-dependencies]
diff --git a/codex-rs/tui/src/conversation_history_widget.rs b/codex-rs/tui/src/conversation_history_widget.rs
index 9242e003..a23e00d7 100644
--- a/codex-rs/tui/src/conversation_history_widget.rs
+++ b/codex-rs/tui/src/conversation_history_widget.rs
@@ -299,7 +299,6 @@ impl ConversationHistoryWidget {
         for entry in self.entries.iter_mut() {
             if let HistoryCell::ActiveMcpToolCall {
                 call_id: history_id,
-                fq_tool_name,
                 invocation,
                 start,
                 ..
@@ -307,7 +306,7 @@ impl ConversationHistoryWidget {
             {
                 if &call_id == history_id {
                     let completed = HistoryCell::new_completed_mcp_tool_call(
-                        fq_tool_name.clone(),
+                        width,
                         invocation.clone(),
                         *start,
                         success,
diff --git a/codex-rs/tui/src/history_cell.rs b/codex-rs/tui/src/history_cell.rs
index a1fc672c..481576b5 100644
--- a/codex-rs/tui/src/history_cell.rs
+++ b/codex-rs/tui/src/history_cell.rs
@@ -2,6 +2,7 @@ use crate::cell_widget::CellWidget;
 use crate::exec_command::escape_command;
 use crate::markdown::append_markdown;
 use crate::text_block::TextBlock;
+use crate::text_formatting::format_and_truncate_tool_result;
 use base64::Engine;
 use codex_ansi_escape::ansi_escape_line;
 use codex_common::elapsed::format_duration;
@@ -14,6 +15,7 @@ use image::DynamicImage;
 use image::GenericImageView;
 use image::ImageReader;
 use lazy_static::lazy_static;
+use mcp_types::EmbeddedResourceResource;
 use ratatui::prelude::*;
 use ratatui::style::Color;
 use ratatui::style::Modifier;
@@ -73,18 +75,14 @@ pub(crate) enum HistoryCell {
     /// An MCP tool call that has not finished yet.
     ActiveMcpToolCall {
         call_id: String,
-        /// `server.tool` fully-qualified name so we can show a concise label
-        fq_tool_name: String,
-        /// Formatted invocation that mirrors the `$ cmd ...` style of exec
-        /// commands. We keep this around so the completed state can reuse the
-        /// exact same text without re-formatting.
-        invocation: String,
+        /// Formatted line that shows the command name and arguments
+        invocation: Line<'static>,
         start: Instant,
         view: TextBlock,
     },
 
     /// Completed MCP tool call where we show the result serialized as JSON.
-    CompletedMcpToolCallWithTextOutput { view: TextBlock },
+    CompletedMcpToolCall { view: TextBlock },
 
     /// Completed MCP tool call where the result is an image.
     /// Admittedly, [mcp_types::CallToolResult] can have multiple content types,
@@ -289,8 +287,6 @@ impl HistoryCell {
         tool: String,
         arguments: Option<serde_json::Value>,
     ) -> Self {
-        let fq_tool_name = format!("{server}.{tool}");
-
         // Format the arguments as compact JSON so they roughly fit on one
         // line. If there are no arguments we keep it empty so the invocation
         // mirrors a function-style call.
@@ -302,29 +298,30 @@ impl HistoryCell {
             })
             .unwrap_or_default();
 
-        let invocation = if args_str.is_empty() {
-            format!("{fq_tool_name}()")
-        } else {
-            format!("{fq_tool_name}({args_str})")
-        };
+        let invocation_spans = vec![
+            Span::styled(server, Style::default().fg(Color::Blue)),
+            Span::raw("."),
+            Span::styled(tool, Style::default().fg(Color::Blue)),
+            Span::raw("("),
+            Span::styled(args_str, Style::default().fg(Color::Gray)),
+            Span::raw(")"),
+        ];
+        let invocation = Line::from(invocation_spans);
 
         let start = Instant::now();
         let title_line = Line::from(vec!["tool".magenta(), " running...".dim()]);
-        let lines: Vec<Line<'static>> = vec![
-            title_line,
-            Line::from(format!("$ {invocation}")),
-            Line::from(""),
-        ];
+        let lines: Vec<Line<'static>> = vec![title_line, invocation.clone(), Line::from("")];
 
         HistoryCell::ActiveMcpToolCall {
             call_id,
-            fq_tool_name,
             invocation,
             start,
             view: TextBlock::new(lines),
         }
     }
 
+    /// If the first content is an image, return a new cell with the image.
+    /// TODO(rgwood-dd): Handle images properly even if they're not the first result.
     fn try_new_completed_mcp_tool_call_with_image_output(
         result: &Result<mcp_types::CallToolResult, String>,
     ) -> Option<Self> {
@@ -370,8 +367,8 @@ impl HistoryCell {
     }
 
     pub(crate) fn new_completed_mcp_tool_call(
-        fq_tool_name: String,
-        invocation: String,
+        num_cols: u16,
+        invocation: Line<'static>,
         start: Instant,
         success: bool,
         result: Result<mcp_types::CallToolResult, String>,
@@ -384,36 +381,70 @@ impl HistoryCell {
         let status_str = if success { "success" } else { "failed" };
         let title_line = Line::from(vec![
             "tool".magenta(),
-            format!(" {fq_tool_name} ({status_str}, duration: {})", duration).dim(),
+            " ".into(),
+            if success {
+                status_str.green()
+            } else {
+                status_str.red()
+            },
+            format!(", duration: {duration}").gray(),
         ]);
 
         let mut lines: Vec<Line<'static>> = Vec::new();
         lines.push(title_line);
-        lines.push(Line::from(format!("$ {invocation}")));
+        lines.push(invocation);
 
-        // Convert result into serde_json::Value early so we don't have to
-        // worry about lifetimes inside the match arm.
-        let result_val = result.map(|r| {
-            serde_json::to_value(r)
-                .unwrap_or_else(|_| serde_json::Value::String("<serialization error>".into()))
-        });
+        match result {
+            Ok(mcp_types::CallToolResult { content, .. }) => {
+                if !content.is_empty() {
+                    lines.push(Line::from(""));
 
-        if let Ok(res_val) = result_val {
-            let json_pretty =
-                serde_json::to_string_pretty(&res_val).unwrap_or_else(|_| res_val.to_string());
-            let mut iter = json_pretty.lines();
-            for raw in iter.by_ref().take(TOOL_CALL_MAX_LINES) {
-                lines.push(Line::from(raw.to_string()).dim());
+                    for tool_call_result in content {
+                        let line_text = match tool_call_result {
+                            mcp_types::CallToolResultContent::TextContent(text) => {
+                                format_and_truncate_tool_result(
+                                    &text.text,
+                                    TOOL_CALL_MAX_LINES,
+                                    num_cols as usize,
+                                )
+                            }
+                            mcp_types::CallToolResultContent::ImageContent(_) => {
+                                // TODO show images even if they're not the first result, will require a refactor of `CompletedMcpToolCall`
+                                "<image content>".to_string()
+                            }
+                            mcp_types::CallToolResultContent::AudioContent(_) => {
+                                "<audio content>".to_string()
+                            }
+                            mcp_types::CallToolResultContent::EmbeddedResource(resource) => {
+                                let uri = match resource.resource {
+                                    EmbeddedResourceResource::TextResourceContents(text) => {
+                                        text.uri
+                                    }
+                                    EmbeddedResourceResource::BlobResourceContents(blob) => {
+                                        blob.uri
+                                    }
+                                };
+                                format!("embedded resource: {uri}")
+                            }
+                        };
+                        lines.push(Line::styled(line_text, Style::default().fg(Color::Gray)));
+                    }
+                }
+
+                lines.push(Line::from(""));
             }
-            let remaining = iter.count();
-            if remaining > 0 {
-                lines.push(Line::from(format!("... {} additional lines", remaining)).dim());
+            Err(e) => {
+                lines.push(Line::from(vec![
+                    Span::styled(
+                        "Error: ",
+                        Style::default().fg(Color::Red).add_modifier(Modifier::BOLD),
+                    ),
+                    Span::raw(e),
+                ]));
             }
-        }
+        };
 
-        lines.push(Line::from(""));
-
-        HistoryCell::CompletedMcpToolCallWithTextOutput {
+        HistoryCell::CompletedMcpToolCall {
             view: TextBlock::new(lines),
         }
     }
@@ -520,7 +551,7 @@ impl CellWidget for HistoryCell {
             | HistoryCell::ErrorEvent { view }
             | HistoryCell::SessionInfo { view }
             | HistoryCell::CompletedExecCommand { view }
-            | HistoryCell::CompletedMcpToolCallWithTextOutput { view }
+            | HistoryCell::CompletedMcpToolCall { view }
             | HistoryCell::PendingPatch { view }
             | HistoryCell::ActiveExecCommand { view, .. }
             | HistoryCell::ActiveMcpToolCall { view, .. } => view.height(width),
@@ -541,7 +572,7 @@ impl CellWidget for HistoryCell {
             | HistoryCell::ErrorEvent { view }
             | HistoryCell::SessionInfo { view }
             | HistoryCell::CompletedExecCommand { view }
-            | HistoryCell::CompletedMcpToolCallWithTextOutput { view }
+            | HistoryCell::CompletedMcpToolCall { view }
             | HistoryCell::PendingPatch { view }
             | HistoryCell::ActiveExecCommand { view, .. }
             | HistoryCell::ActiveMcpToolCall { view, .. } => {
diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs
index df85673e..535168e3 100644
--- a/codex-rs/tui/src/lib.rs
+++ b/codex-rs/tui/src/lib.rs
@@ -34,6 +34,7 @@ mod scroll_event_helper;
 mod slash_command;
 mod status_indicator_widget;
 mod text_block;
+mod text_formatting;
 mod tui;
 mod user_approval_widget;
 
diff --git a/codex-rs/tui/src/text_formatting.rs b/codex-rs/tui/src/text_formatting.rs
new file mode 100644
index 00000000..e79fbf7a
--- /dev/null
+++ b/codex-rs/tui/src/text_formatting.rs
@@ -0,0 +1,268 @@
+use unicode_segmentation::UnicodeSegmentation;
+
+/// Truncate a tool result to fit within the given height and width. If the text is valid JSON, we format it in a compact way before truncating.
+/// This is a best-effort approach that may not work perfectly for text where 1 grapheme is rendered as multiple terminal cells.
+pub(crate) fn format_and_truncate_tool_result(
+    text: &str,
+    max_lines: usize,
+    line_width: usize,
+) -> String {
+    // Work out the maximum number of graphemes we can display for a result.
+    // It's not guaranteed that 1 grapheme = 1 cell, so we subtract 1 per line as a fudge factor.
+    // It also won't handle future terminal resizes properly, but it's an OK approximation for now.
+    let max_graphemes = (max_lines * line_width).saturating_sub(max_lines);
+
+    if let Some(formatted_json) = format_json_compact(text) {
+        truncate_text(&formatted_json, max_graphemes)
+    } else {
+        truncate_text(text, max_graphemes)
+    }
+}
+
+/// Format JSON text in a compact single-line format with spaces for better Ratatui wrapping.
+/// Ex: `{"a":"b",c:["d","e"]}` -> `{"a": "b", "c": ["d", "e"]}`
+/// Returns the formatted JSON string if the input is valid JSON, otherwise returns None.
+/// This is a little complicated, but it's necessary because Ratatui's wrapping is *very* limited
+/// and can only do line breaks at whitespace. If we use the default serde_json format, we get lines
+/// without spaces that Ratatui can't wrap nicely. If we use the serde_json pretty format as-is,
+/// it's much too sparse and uses too many terminal rows.
+/// Relevant issue: https://github.com/ratatui/ratatui/issues/293
+pub(crate) fn format_json_compact(text: &str) -> Option<String> {
+    let json = serde_json::from_str::<serde_json::Value>(text).ok()?;
+    let json_pretty = serde_json::to_string_pretty(&json).unwrap_or_else(|_| json.to_string());
+
+    // Convert multi-line pretty JSON to compact single-line format by removing newlines and excess whitespace
+    let mut result = String::new();
+    let mut chars = json_pretty.chars().peekable();
+    let mut in_string = false;
+    let mut escape_next = false;
+
+    // Iterate over the characters in the JSON string, adding spaces after : and , but only when not in a string
+    while let Some(ch) = chars.next() {
+        match ch {
+            '"' if !escape_next => {
+                in_string = !in_string;
+                result.push(ch);
+            }
+            '\\' if in_string => {
+                escape_next = !escape_next;
+                result.push(ch);
+            }
+            '\n' | '\r' if !in_string => {
+                // Skip newlines when not in a string
+            }
+            ' ' | '\t' if !in_string => {
+                // Add a space after : and , but only when not in a string
+                if let Some(&next_ch) = chars.peek() {
+                    if let Some(last_ch) = result.chars().last() {
+                        if (last_ch == ':' || last_ch == ',') && !matches!(next_ch, '}' | ']') {
+                            result.push(' ');
+                        }
+                    }
+                }
+            }
+            _ => {
+                if escape_next && in_string {
+                    escape_next = false;
+                }
+                result.push(ch);
+            }
+        }
+    }
+
+    Some(result)
+}
+
+/// Truncate `text` to `max_graphemes` graphemes. Using graphemes to avoid accidentally truncating in the middle of a multi-codepoint character.
+pub(crate) fn truncate_text(text: &str, max_graphemes: usize) -> String {
+    let mut graphemes = text.grapheme_indices(true);
+
+    // Check if there's a grapheme at position max_graphemes (meaning there are more than max_graphemes total)
+    if let Some((byte_index, _)) = graphemes.nth(max_graphemes) {
+        // There are more than max_graphemes, so we need to truncate
+        if max_graphemes >= 3 {
+            // Truncate to max_graphemes - 3 and add "..." to stay within limit
+            let mut truncate_graphemes = text.grapheme_indices(true);
+            if let Some((truncate_byte_index, _)) = truncate_graphemes.nth(max_graphemes - 3) {
+                let truncated = &text[..truncate_byte_index];
+                format!("{}...", truncated)
+            } else {
+                text.to_string()
+            }
+        } else {
+            // max_graphemes < 3, so just return first max_graphemes without "..."
+            let truncated = &text[..byte_index];
+            truncated.to_string()
+        }
+    } else {
+        // There are max_graphemes or fewer graphemes, return original text
+        text.to_string()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #![allow(clippy::unwrap_used)]
+    use super::*;
+    use pretty_assertions::assert_eq;
+
+    #[test]
+    fn test_truncate_text() {
+        let text = "Hello, world!";
+        let truncated = truncate_text(text, 8);
+        assert_eq!(truncated, "Hello...");
+    }
+
+    #[test]
+    fn test_truncate_empty_string() {
+        let text = "";
+        let truncated = truncate_text(text, 5);
+        assert_eq!(truncated, "");
+    }
+
+    #[test]
+    fn test_truncate_max_graphemes_zero() {
+        let text = "Hello";
+        let truncated = truncate_text(text, 0);
+        assert_eq!(truncated, "");
+    }
+
+    #[test]
+    fn test_truncate_max_graphemes_one() {
+        let text = "Hello";
+        let truncated = truncate_text(text, 1);
+        assert_eq!(truncated, "H");
+    }
+
+    #[test]
+    fn test_truncate_max_graphemes_two() {
+        let text = "Hello";
+        let truncated = truncate_text(text, 2);
+        assert_eq!(truncated, "He");
+    }
+
+    #[test]
+    fn test_truncate_max_graphemes_three_boundary() {
+        let text = "Hello";
+        let truncated = truncate_text(text, 3);
+        assert_eq!(truncated, "...");
+    }
+
+    #[test]
+    fn test_truncate_text_shorter_than_limit() {
+        let text = "Hi";
+        let truncated = truncate_text(text, 10);
+        assert_eq!(truncated, "Hi");
+    }
+
+    #[test]
+    fn test_truncate_text_exact_length() {
+        let text = "Hello";
+        let truncated = truncate_text(text, 5);
+        assert_eq!(truncated, "Hello");
+    }
+
+    #[test]
+    fn test_truncate_emoji() {
+        let text = "👋🌍🚀✨💫";
+        let truncated = truncate_text(text, 3);
+        assert_eq!(truncated, "...");
+
+        let truncated_longer = truncate_text(text, 4);
+        assert_eq!(truncated_longer, "👋...");
+    }
+
+    #[test]
+    fn test_truncate_unicode_combining_characters() {
+        let text = "é́ñ̃"; // Characters with combining marks
+        let truncated = truncate_text(text, 2);
+        assert_eq!(truncated, "é́ñ̃");
+    }
+
+    #[test]
+    fn test_truncate_very_long_text() {
+        let text = "a".repeat(1000);
+        let truncated = truncate_text(&text, 10);
+        assert_eq!(truncated, "aaaaaaa...");
+        assert_eq!(truncated.len(), 10); // 7 'a's + 3 dots
+    }
+
+    #[test]
+    fn test_format_json_compact_simple_object() {
+        let json = r#"{ "name": "John", "age": 30 }"#;
+        let result = format_json_compact(json).unwrap();
+        assert_eq!(result, r#"{"name": "John", "age": 30}"#);
+    }
+
+    #[test]
+    fn test_format_json_compact_nested_object() {
+        let json = r#"{ "user": { "name": "John", "details": { "age": 30, "city": "NYC" } } }"#;
+        let result = format_json_compact(json).unwrap();
+        assert_eq!(
+            result,
+            r#"{"user": {"name": "John", "details": {"age": 30, "city": "NYC"}}}"#
+        );
+    }
+
+    #[test]
+    fn test_format_json_compact_array() {
+        let json = r#"[ 1, 2, { "key": "value" }, "string" ]"#;
+        let result = format_json_compact(json).unwrap();
+        assert_eq!(result, r#"[1, 2, {"key": "value"}, "string"]"#);
+    }
+
+    #[test]
+    fn test_format_json_compact_already_compact() {
+        let json = r#"{"compact":true}"#;
+        let result = format_json_compact(json).unwrap();
+        assert_eq!(result, r#"{"compact": true}"#);
+    }
+
+    #[test]
+    fn test_format_json_compact_with_whitespace() {
+        let json = r#"
+        {
+            "name": "John",
+            "hobbies": [
+                "reading",
+                "coding"
+            ]
+        }
+        "#;
+        let result = format_json_compact(json).unwrap();
+        assert_eq!(
+            result,
+            r#"{"name": "John", "hobbies": ["reading", "coding"]}"#
+        );
+    }
+
+    #[test]
+    fn test_format_json_compact_invalid_json() {
+        let invalid_json = r#"{"invalid": json syntax}"#;
+        let result = format_json_compact(invalid_json);
+        assert!(result.is_none());
+    }
+
+    #[test]
+    fn test_format_json_compact_empty_object() {
+        let json = r#"{}"#;
+        let result = format_json_compact(json).unwrap();
+        assert_eq!(result, "{}");
+    }
+
+    #[test]
+    fn test_format_json_compact_empty_array() {
+        let json = r#"[]"#;
+        let result = format_json_compact(json).unwrap();
+        assert_eq!(result, "[]");
+    }
+
+    #[test]
+    fn test_format_json_compact_primitive_values() {
+        assert_eq!(format_json_compact("42").unwrap(), "42");
+        assert_eq!(format_json_compact("true").unwrap(), "true");
+        assert_eq!(format_json_compact("false").unwrap(), "false");
+        assert_eq!(format_json_compact("null").unwrap(), "null");
+        assert_eq!(format_json_compact(r#""string""#).unwrap(), r#""string""#);
+    }
+}