diff --git a/crates/agent/src/active_thread.rs b/crates/agent/src/active_thread.rs
index 9182aa0c8a..0ca4033b1f 100644
--- a/crates/agent/src/active_thread.rs
+++ b/crates/agent/src/active_thread.rs
@@ -327,6 +327,7 @@ fn tool_use_markdown_style(window: &Window, cx: &mut App) -> MarkdownStyle {
     }
 }
 
+const CODEBLOCK_CONTAINER_GROUP: &str = "codeblock_container";
 const MAX_UNCOLLAPSED_LINES_IN_CODE_BLOCK: usize = 10;
 
 fn render_markdown_code_block(
@@ -485,12 +486,18 @@ fn render_markdown_code_block(
         .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(true);
+    let can_expand = metadata.line_count > MAX_UNCOLLAPSED_LINES_IN_CODE_BLOCK;
+
+    let is_expanded = if can_expand {
+        active_thread
+            .read(cx)
+            .expanded_code_blocks
+            .get(&(message_id, ix))
+            .copied()
+            .unwrap_or(false)
+    } else {
+        false
+    };
 
     let codeblock_header_bg = cx
         .theme()
@@ -511,7 +518,7 @@ fn render_markdown_code_block(
         .children(label)
         .child(
             h_flex()
-                .visible_on_hover("codeblock_container")
+                .visible_on_hover(CODEBLOCK_CONTAINER_GROUP)
                 .gap_1()
                 .child(
                     IconButton::new(
@@ -553,45 +560,42 @@ fn render_markdown_code_block(
                         }
                     }),
                 )
-                .when(
-                    metadata.line_count > MAX_UNCOLLAPSED_LINES_IN_CODE_BLOCK,
-                    |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"
+                .when(can_expand, |header| {
+                    header.child(
+                        IconButton::new(
+                            ("expand-collapse-code", ix),
+                            if is_expanded {
+                                IconName::ChevronUp
                             } 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(true);
-                                        *is_expanded = !*is_expanded;
-                                        cx.notify();
-                                    });
-                                }
-                            }),
+                                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(true);
+                                    *is_expanded = !*is_expanded;
+                                    cx.notify();
+                                });
+                            }
+                        }),
+                    )
+                }),
         );
 
     v_flex()
-        .group("codeblock_container")
+        .group(CODEBLOCK_CONTAINER_GROUP)
         .my_2()
         .overflow_hidden()
         .rounded_lg()
@@ -599,10 +603,7 @@ fn render_markdown_code_block(
         .border_color(cx.theme().colors().border.opacity(0.6))
         .bg(cx.theme().colors().editor_background)
         .child(codeblock_header)
-        .when(
-            metadata.line_count > MAX_UNCOLLAPSED_LINES_IN_CODE_BLOCK && !is_expanded,
-            |this| this.max_h_80(),
-        )
+        .when(can_expand && !is_expanded, |this| this.max_h_80())
 }
 
 fn render_code_language(
@@ -2353,19 +2354,22 @@ impl ActiveThread {
                                             let editor_bg = cx.theme().colors().editor_background;
 
                                             move |el, range, metadata, _, cx| {
+                                                let can_expand = metadata.line_count
+                                                    > MAX_UNCOLLAPSED_LINES_IN_CODE_BLOCK;
+                                                if !can_expand {
+                                                    return el;
+                                                }
+
                                                 let is_expanded = active_thread
                                                     .read(cx)
                                                     .expanded_code_blocks
                                                     .get(&(message_id, range.start))
                                                     .copied()
-                                                    .unwrap_or(true);
-
-                                                if is_expanded
-                                                    || metadata.line_count
-                                                        <= MAX_UNCOLLAPSED_LINES_IN_CODE_BLOCK
-                                                {
+                                                    .unwrap_or(false);
+                                                if is_expanded {
                                                     return el;
                                                 }
+
                                                 el.child(
                                                     div()
                                                         .absolute()