agent: Collapse code blocks in the active thread (#28467)

Danilo Leal and Bennet Bo Fenner created

Release Notes:

- N/A

---------

Co-authored-by: Bennet Bo Fenner <bennetbo@gmx.de>

Change summary

crates/agent/src/active_thread.rs | 210 ++++++++++++++++++++++++--------
crates/markdown/src/markdown.rs   |  46 ++++--
2 files changed, 187 insertions(+), 69 deletions(-)

Detailed changes

crates/agent/src/active_thread.rs 🔗

@@ -57,6 +57,7 @@ pub struct ActiveThread {
     editing_message: Option<(MessageId, EditMessageState)>,
     expanded_tool_uses: HashMap<LanguageModelToolUseId, bool>,
     expanded_thinking_segments: HashMap<(MessageId, usize), bool>,
+    expanded_code_blocks: HashMap<(MessageId, usize), bool>,
     last_error: Option<ThreadError>,
     notifications: Vec<WindowHandle<AgentNotification>>,
     copied_code_block_ids: HashSet<(MessageId, usize)>,
@@ -297,7 +298,7 @@ fn render_markdown_code_block(
     codeblock_range: Range<usize>,
     active_thread: Entity<ActiveThread>,
     workspace: WeakEntity<Workspace>,
-    _window: &mut Window,
+    _window: &Window,
     cx: &App,
 ) -> Div {
     let label = match kind {
@@ -377,16 +378,20 @@ fn render_markdown_code_block(
                 .rounded_sm()
                 .hover(|item| item.bg(cx.theme().colors().element_hover.opacity(0.5)))
                 .tooltip(Tooltip::text("Jump to File"))
-                .children(
-                    file_icons::FileIcons::get_icon(&path_range.path, cx)
-                        .map(Icon::from_path)
-                        .map(|icon| icon.color(Color::Muted).size(IconSize::XSmall)),
-                )
-                .child(content)
                 .child(
-                    Icon::new(IconName::ArrowUpRight)
-                        .size(IconSize::XSmall)
-                        .color(Color::Ignored),
+                    h_flex()
+                        .gap_0p5()
+                        .children(
+                            file_icons::FileIcons::get_icon(&path_range.path, cx)
+                                .map(Icon::from_path)
+                                .map(|icon| icon.color(Color::Muted).size(IconSize::XSmall)),
+                        )
+                        .child(content)
+                        .child(
+                            Icon::new(IconName::ArrowUpRight)
+                                .size(IconSize::XSmall)
+                                .color(Color::Ignored),
+                        ),
                 )
                 .on_click({
                     let path_range = path_range.clone();
@@ -444,16 +449,29 @@ fn render_markdown_code_block(
         }),
     };
 
+    let codeblock_was_copied = active_thread
+        .read(cx)
+        .copied_code_block_ids
+        .contains(&(message_id, ix));
+
+    let is_expanded = active_thread
+        .read(cx)
+        .expanded_code_blocks
+        .get(&(message_id, ix))
+        .copied()
+        .unwrap_or(false);
+
     let codeblock_header_bg = cx
         .theme()
         .colors()
         .element_background
         .blend(cx.theme().colors().editor_foreground.opacity(0.01));
 
-    let codeblock_was_copied = active_thread
-        .read(cx)
-        .copied_code_block_ids
-        .contains(&(message_id, ix));
+    let line_count = without_fences(&parsed_markdown.source()[codeblock_range.clone()])
+        .lines()
+        .count();
+
+    const MAX_COLLAPSED_LINES: usize = 5;
 
     let codeblock_header = h_flex()
         .group("codeblock_header")
@@ -466,57 +484,104 @@ fn render_markdown_code_block(
         .rounded_t_md()
         .children(label)
         .child(
-            div().visible_on_hover("codeblock_header").child(
-                IconButton::new(
-                    ("copy-markdown-code", ix),
-                    if codeblock_was_copied {
-                        IconName::Check
-                    } else {
-                        IconName::Copy
-                    },
-                )
-                .icon_color(Color::Muted)
-                .shape(ui::IconButtonShape::Square)
-                .tooltip(Tooltip::text("Copy Code"))
-                .on_click({
-                    let active_thread = active_thread.clone();
-                    let parsed_markdown = parsed_markdown.clone();
-                    move |_event, _window, cx| {
-                        active_thread.update(cx, |this, cx| {
-                            this.copied_code_block_ids.insert((message_id, ix));
-
-                            let code =
-                                without_fences(&parsed_markdown.source()[codeblock_range.clone()])
+            h_flex()
+                .gap_1()
+                .child(
+                    div().visible_on_hover("codeblock_header").child(
+                        IconButton::new(
+                            ("copy-markdown-code", ix),
+                            if codeblock_was_copied {
+                                IconName::Check
+                            } else {
+                                IconName::Copy
+                            },
+                        )
+                        .icon_color(Color::Muted)
+                        .shape(ui::IconButtonShape::Square)
+                        .tooltip(Tooltip::text("Copy Code"))
+                        .on_click({
+                            let active_thread = active_thread.clone();
+                            let parsed_markdown = parsed_markdown.clone();
+                            move |_event, _window, cx| {
+                                active_thread.update(cx, |this, cx| {
+                                    this.copied_code_block_ids.insert((message_id, ix));
+
+                                    let code = without_fences(
+                                        &parsed_markdown.source()[codeblock_range.clone()],
+                                    )
                                     .to_string();
 
-                            cx.write_to_clipboard(ClipboardItem::new_string(code.clone()));
+                                    cx.write_to_clipboard(ClipboardItem::new_string(code.clone()));
 
-                            cx.spawn(async move |this, cx| {
-                                cx.background_executor().timer(Duration::from_secs(2)).await;
+                                    cx.spawn(async move |this, cx| {
+                                        cx.background_executor()
+                                            .timer(Duration::from_secs(2))
+                                            .await;
 
-                                cx.update(|cx| {
-                                    this.update(cx, |this, cx| {
-                                        this.copied_code_block_ids.remove(&(message_id, ix));
-                                        cx.notify();
+                                        cx.update(|cx| {
+                                            this.update(cx, |this, cx| {
+                                                this.copied_code_block_ids
+                                                    .remove(&(message_id, ix));
+                                                cx.notify();
+                                            })
+                                        })
+                                        .ok();
                                     })
-                                })
-                                .ok();
-                            })
-                            .detach();
-                        });
-                    }
+                                    .detach();
+                                });
+                            }
+                        }),
+                    ),
+                )
+                .when(line_count > MAX_COLLAPSED_LINES, |header| {
+                    header.child(
+                        IconButton::new(
+                            ("expand-collapse-code", ix),
+                            if is_expanded {
+                                IconName::ChevronUp
+                            } else {
+                                IconName::ChevronDown
+                            },
+                        )
+                        .icon_color(Color::Muted)
+                        .shape(ui::IconButtonShape::Square)
+                        .tooltip(Tooltip::text(if is_expanded {
+                            "Collapse Code"
+                        } else {
+                            "Expand Code"
+                        }))
+                        .on_click({
+                            let active_thread = active_thread.clone();
+                            move |_event, _window, cx| {
+                                active_thread.update(cx, |this, cx| {
+                                    let is_expanded = this
+                                        .expanded_code_blocks
+                                        .entry((message_id, ix))
+                                        .or_insert(false);
+                                    *is_expanded = !*is_expanded;
+                                    cx.notify();
+                                });
+                            }
+                        }),
+                    )
                 }),
-            ),
         );
 
     v_flex()
-        .mb_2()
-        .relative()
+        .my_2()
         .overflow_hidden()
         .rounded_lg()
         .border_1()
         .border_color(cx.theme().colors().border_variant)
+        .bg(cx.theme().colors().editor_background)
         .child(codeblock_header)
+        .when(line_count > MAX_COLLAPSED_LINES, |this| {
+            if is_expanded {
+                this.h_full()
+            } else {
+                this.max_h_40()
+            }
+        })
 }
 
 fn open_markdown_link(
@@ -626,6 +691,7 @@ impl ActiveThread {
             rendered_tool_uses: HashMap::default(),
             expanded_tool_uses: HashMap::default(),
             expanded_thinking_segments: HashMap::default(),
+            expanded_code_blocks: HashMap::default(),
             list_state: list_state.clone(),
             scrollbar_state: ScrollbarState::new(list_state),
             show_scrollbar: false,
@@ -1835,10 +1901,10 @@ impl ActiveThread {
                                     render: Arc::new({
                                         let workspace = workspace.clone();
                                         let active_thread = cx.entity();
-                                        move |id, kind, parsed_markdown, range, window, cx| {
+                                        move |kind, parsed_markdown, range, window, cx| {
                                             render_markdown_code_block(
                                                 message_id,
-                                                id,
+                                                range.start,
                                                 kind,
                                                 parsed_markdown,
                                                 range,
@@ -1849,6 +1915,44 @@ impl ActiveThread {
                                             )
                                         }
                                     }),
+                                    transform: Some(Arc::new({
+                                        let active_thread = cx.entity();
+                                        move |el, range, _, cx| {
+                                            let is_expanded = active_thread
+                                                .read(cx)
+                                                .expanded_code_blocks
+                                                .get(&(message_id, range.start))
+                                                .copied()
+                                                .unwrap_or(false);
+
+                                            if is_expanded {
+                                                return el;
+                                            }
+                                            el.child(
+                                                div()
+                                                    .absolute()
+                                                    .bottom_0()
+                                                    .left_0()
+                                                    .w_full()
+                                                    .h_1_4()
+                                                    .rounded_b_lg()
+                                                    .bg(gpui::linear_gradient(
+                                                        0.,
+                                                        gpui::linear_color_stop(
+                                                            cx.theme().colors().editor_background,
+                                                            0.,
+                                                        ),
+                                                        gpui::linear_color_stop(
+                                                            cx.theme()
+                                                                .colors()
+                                                                .editor_background
+                                                                .opacity(0.),
+                                                            1.,
+                                                        ),
+                                                    )),
+                                            )
+                                        }
+                                    })),
                                 })
                                 .on_url_click({
                                     let workspace = self.workspace.clone();

crates/markdown/src/markdown.rs 🔗

@@ -88,12 +88,21 @@ struct Options {
 }
 
 pub enum CodeBlockRenderer {
-    Default { copy_button: bool },
-    Custom { render: CodeBlockRenderFn },
+    Default {
+        copy_button: bool,
+    },
+    Custom {
+        render: CodeBlockRenderFn,
+        /// A function that can modify the parent container after the code block
+        /// content has been appended as a child element.
+        transform: Option<CodeBlockTransformFn>,
+    },
 }
 
 pub type CodeBlockRenderFn =
-    Arc<dyn Fn(usize, &CodeBlockKind, &ParsedMarkdown, Range<usize>, &mut Window, &App) -> Div>;
+    Arc<dyn Fn(&CodeBlockKind, &ParsedMarkdown, Range<usize>, &mut Window, &App) -> Div>;
+
+pub type CodeBlockTransformFn = Arc<dyn Fn(AnyDiv, Range<usize>, &mut Window, &App) -> AnyDiv>;
 
 actions!(markdown, [Copy, CopyAsMarkdown]);
 
@@ -594,7 +603,7 @@ impl Element for MarkdownElement {
             0
         };
 
-        for (index, (range, event)) in parsed_markdown.events.iter().enumerate() {
+        for (range, event) in parsed_markdown.events.iter() {
             match event {
                 MarkdownEvent::Start(tag) => {
                     match tag {
@@ -676,15 +685,9 @@ impl Element for MarkdownElement {
                                     builder.push_code_block(language);
                                     builder.push_div(code_block, range, markdown_end);
                                 }
-                                (CodeBlockRenderer::Custom { render }, _) => {
-                                    let parent_container = render(
-                                        index,
-                                        kind,
-                                        &parsed_markdown,
-                                        range.clone(),
-                                        window,
-                                        cx,
-                                    );
+                                (CodeBlockRenderer::Custom { render, .. }, _) => {
+                                    let parent_container =
+                                        render(kind, &parsed_markdown, range.clone(), window, cx);
 
                                     builder.push_div(parent_container, range, markdown_end);
 
@@ -695,9 +698,12 @@ impl Element for MarkdownElement {
                                             if self.style.code_block_overflow_x_scroll {
                                                 code_block.style().restrict_scroll_to_axis =
                                                     Some(true);
-                                                code_block.flex().overflow_x_scroll()
+                                                code_block
+                                                    .flex()
+                                                    .overflow_x_scroll()
+                                                    .overflow_y_hidden()
                                             } else {
-                                                code_block.w_full()
+                                                code_block.w_full().overflow_hidden()
                                             }
                                         });
 
@@ -846,6 +852,14 @@ impl Element for MarkdownElement {
                             builder.pop_text_style();
                         }
 
+                        if let CodeBlockRenderer::Custom {
+                            transform: Some(modify),
+                            ..
+                        } = &self.code_block_renderer
+                        {
+                            builder.modify_current_div(|el| modify(el, range.clone(), window, cx));
+                        }
+
                         if matches!(
                             &self.code_block_renderer,
                             CodeBlockRenderer::Default { copy_button: true }
@@ -1049,7 +1063,7 @@ impl IntoElement for MarkdownElement {
     }
 }
 
-enum AnyDiv {
+pub enum AnyDiv {
     Div(Div),
     Stateful(Stateful<Div>),
 }