TUI: remeasure running agent commands on row growth

Co-Authored-By: Oz <oz-agent@warp.dev>

Author: harryalbert

crates/warp_tui/src/agent_block.rs

@@ -1126,16 +1126,13 @@ impl TuiAIBlock {
     /// Whether the cached height is stale at `width`.
     pub(super) fn needs_height_measurement(&self, width: u16, app: &AppContext) -> bool {
         self.last_measured_width.get() != Some(width)
-            || self.block_model.status(app).is_streaming()
             || self.action_views.values().any(|view| match view {
                 TuiToolCallView::AskQuestion(_)
                 | TuiToolCallView::FileEdits(_)
                 | TuiToolCallView::Generic(_)
                 | TuiToolCallView::Plan(_)
                 | TuiToolCallView::OrchestrationBlock(_) => false,
-                TuiToolCallView::ShellCommand(view) => {
-                    view.as_ref(app).needs_continuous_height_measurement()
-                }
+                TuiToolCallView::ShellCommand(view) => view.as_ref(app).needs_height_measurement(),
             })
     }
 
@@ -1178,16 +1175,35 @@ impl TuiAIBlock {
         ctx: &mut TuiLayoutContext,
         app: &AppContext,
     ) -> usize {
+        let command_extents = self
+            .action_views
+            .values()
+            .filter_map(|view| match view {
+                TuiToolCallView::ShellCommand(view) => {
+                    Some((view, view.as_ref(app).dynamic_content_extent()))
+                }
+                TuiToolCallView::AskQuestion(_)
+                | TuiToolCallView::FileEdits(_)
+                | TuiToolCallView::Generic(_)
+                | TuiToolCallView::Plan(_)
+                | TuiToolCallView::OrchestrationBlock(_) => None,
+            })
+            .collect::<Vec<_>>();
         let mut element = self.render_element(app);
-        usize::from(
+        let height = usize::from(
             element
                 .layout(
                     TuiConstraint::loose(TuiSize::new(width, u16::MAX)),
                     ctx,
                     app,
                 )
                 .height,
-        )
+        );
+        for (view, command_extent) in command_extents {
+            view.as_ref(app)
+                .record_content_extent_measurement(command_extent);
+        }
+        height
     }
 
     /// Logical (unwrapped) text for a selection over this block's text

crates/warp_tui/src/tui_block_list_viewport_source.rs

@@ -103,64 +103,75 @@ impl TuiBlockListViewportSource {
     ///
     /// A non-dirty band block is re-measured only when its cached height cannot
     /// be trusted: its last measurement was at a different width (reflow), it
-    /// has never been measured (no recorded width), or it is still streaming
-    /// (its height can grow without a per-update invalidation — e.g. an
-    /// expanded, still-running shell command). At a stable width with no
-    /// dynamic height, nothing extra is measured and the cached
-    /// `last_laid_out_height` is reused. Off-band blocks keep their cached
-    /// height until they scroll into the band.
+    /// has never been measured (no recorded width), or an expanded running
+    /// shell command's terminal row extent changed. Agent output updates dirty
+    /// their block directly, so streaming status alone does not require layout
+    /// polling. At a stable width with no dynamic height change, the cached
+    /// height is reused. Off-band blocks keep their cached height until they
+    /// scroll into the band.
     fn agent_heights_to_measure(
         &self,
         window: TuiViewportWindow,
         available_width: u16,
         app: &AppContext,
     ) -> HashSet<EntityId> {
-        let mut model = self.model.lock();
-        let mut view_ids = model.block_list_mut().take_dirty_rich_content_items();
+        let mut view_ids = self
+            .model
+            .lock()
+            .block_list_mut()
+            .take_dirty_rich_content_items();
+        let band_view_ids = {
+            let model = self.model.lock();
+            let block_list = model.block_list();
+            let band_top = window.scroll_top.saturating_sub(OVERHANG_ROWS);
+            let band_bottom = window
+                .scroll_top
+                .saturating_add(usize::from(window.viewport_height))
+                .saturating_add(OVERHANG_ROWS);
+            let mut cursor = block_list
+                .block_heights()
+                .cursor::<BlockHeight, BlockHeightSummary>();
+            cursor.seek_clamped(&BlockHeight::from(band_top as f64), SeekBias::Left);
+            let mut band_view_ids = Vec::new();
+            while let Some(item) = cursor.item() {
+                let item_top = cursor.start().height.as_f64().floor().max(0.0) as usize;
+                if item_top >= band_bottom {
+                    break;
+                }
+                let item_bottom = item_top.saturating_add(item.height().as_f64().ceil() as usize);
+                if item_bottom > band_top
+                    && let BlockHeightItem::RichContent(rich_content) = item
+                    && !rich_content.should_hide
+                {
+                    band_view_ids.push(rich_content.view_id);
+                }
+                cursor.next();
+            }
+            band_view_ids
+        };
 
         let agent_blocks = self.agent_blocks.borrow();
         let cli_subagent_blocks = self.cli_subagent_blocks.borrow();
         let handoff_blocks = self.handoff_blocks.borrow();
-        let block_list = model.block_list();
-        let band_top = window.scroll_top.saturating_sub(OVERHANG_ROWS);
-        let band_bottom = window
-            .scroll_top
-            .saturating_add(usize::from(window.viewport_height))
-            .saturating_add(OVERHANG_ROWS);
-        let mut cursor = block_list
-            .block_heights()
-            .cursor::<BlockHeight, BlockHeightSummary>();
-        cursor.seek_clamped(&BlockHeight::from(band_top as f64), SeekBias::Left);
-        while let Some(item) = cursor.item() {
-            let item_top = cursor.start().height.as_f64().floor().max(0.0) as usize;
-            if item_top >= band_bottom {
-                break;
-            }
-            let item_bottom = item_top.saturating_add(item.height().as_f64().ceil() as usize);
-            if item_bottom > band_top
-                && let BlockHeightItem::RichContent(rich_content) = item
-                && !rich_content.should_hide
-            {
-                if let Some(view) = agent_blocks.get(&rich_content.view_id) {
-                    if view
-                        .as_ref(app)
-                        .needs_height_measurement(available_width, app)
-                    {
-                        view_ids.insert(rich_content.view_id);
-                    }
-                } else if let Some(view) = cli_subagent_blocks.get(&rich_content.view_id)
-                    && view
-                        .as_ref(app)
-                        .needs_height_measurement(available_width, app)
+        for view_id in band_view_ids {
+            if let Some(view) = agent_blocks.get(&view_id) {
+                if view
+                    .as_ref(app)
+                    .needs_height_measurement(available_width, app)
                 {
-                    view_ids.insert(rich_content.view_id);
-                } else if let Some(view) = handoff_blocks.get(&rich_content.view_id)
-                    && view.as_ref(app).needs_height_measurement(available_width)
-                {
-                    view_ids.insert(rich_content.view_id);
+                    view_ids.insert(view_id);
                 }
+            } else if let Some(view) = cli_subagent_blocks.get(&view_id)
+                && view
+                    .as_ref(app)
+                    .needs_height_measurement(available_width, app)
+            {
+                view_ids.insert(view_id);
+            } else if let Some(view) = handoff_blocks.get(&view_id)
+                && view.as_ref(app).needs_height_measurement(available_width)
+            {
+                view_ids.insert(view_id);
             }
-            cursor.next();
         }
         view_ids
     }

crates/warp_tui/src/tui_block_list_viewport_source_tests.rs

@@ -363,21 +363,21 @@ fn tui_transcript_scroll_reuses_cached_heights_at_stable_width() {
 }
 
 #[test]
-fn tui_agent_streaming_block_remeasured_at_stable_width() {
+fn tui_agent_streaming_block_reuses_cached_height_without_an_update() {
     App::test((), |mut app| async move {
         app.add_singleton_model(|_| Appearance::mock());
-        // A streaming block's height can grow without a per-update
-        // invalidation, so it must be re-measured at a stable width.
+        // Agent output updates explicitly dirty their rich-content item, so
+        // streaming status alone does not require polling its full layout.
         let (source, model, agent_block) = streaming_agent_block_source(&mut app);
 
         request_top_window(&app, &source, 10);
         source.take_selection_row_resizes();
 
-        // Seed a wrong height at the same width without dirtying; the streaming
-        // block is still re-measured, correcting it.
+        // Seed a wrong height at the same width without dirtying. With no
+        // output update, the cached height remains untouched.
         seed_clean_height(&app, &model, &agent_block, 1234.0, 80);
         request_top_window(&app, &source, 10);
-        assert_ne!(rich_content_height(&model, agent_block.id()), Some(1234.0));
+        assert_eq!(rich_content_height(&model, agent_block.id()), Some(1234.0));
     });
 }
 

crates/warp_tui/src/tui_shell_command_view.rs

@@ -29,7 +29,7 @@ use warpui_core::{
 use crate::agent_block_sections::render_fallback_tool_call_section;
 use crate::editor_view::{TuiEditorView, TuiEditorViewEvent};
 use crate::keybindings::{TUI_BINDING_GROUP, is_tui_owned_binding};
-use crate::terminal_block::TerminalBlockElement;
+use crate::terminal_block::{TerminalBlockElement, block_content_rows};
 use crate::terminal_use::user_controls_running_command;
 use crate::tool_call_labels::{
     CommandBlockState, ResolvedCommandBlock, styled_tool_call_label_spans, tool_call_display_state,
@@ -75,6 +75,12 @@ pub(crate) fn init(app: &mut AppContext) {
     ]);
     app.register_tui_binding_validator::<TuiShellCommandView>(is_tui_owned_binding);
 }
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub(super) struct ShellCommandContentExtent {
+    rendered_height: usize,
+    command_rows: usize,
+    output_rows: usize,
+}
 
 struct ShellCommandViewState {
     collapsed: bool,
@@ -120,7 +126,7 @@ pub(super) struct TuiShellCommandView {
     permission_prompt: ViewHandle<TuiPermissionPrompt>,
     command_was_edited: bool,
     state: ShellCommandViewState,
-    command_running: Cell<bool>,
+    last_measured_content_extent: Cell<Option<ShellCommandContentExtent>>,
     header_mouse_state: MouseStateHandle,
     cli_subagent_view: Option<ViewHandle<TuiCLISubagentView>>,
 }
@@ -196,7 +202,7 @@ impl TuiShellCommandView {
             permission_prompt,
             command_was_edited: false,
             state: ShellCommandViewState::new_collapsed(),
-            command_running: Cell::new(false),
+            last_measured_content_extent: Cell::new(None),
             header_mouse_state: MouseStateHandle::default(),
             cli_subagent_view: None,
         }
@@ -343,9 +349,31 @@ impl TuiShellCommandView {
             }
         });
     }
-    /// Whether expanded command output can still grow between layout events.
-    pub(super) fn needs_continuous_height_measurement(&self) -> bool {
-        !self.state.is_collapsed() && self.command_running.get()
+    /// Returns the live row extent when command output can change this view's height.
+    pub(super) fn dynamic_content_extent(&self) -> Option<ShellCommandContentExtent> {
+        let model = self.terminal_model.lock();
+        let block = model.block_list().block_for_ai_action_id(&self.action.id)?;
+        let expanded = !self.state.is_collapsed() || user_controls_running_command(block);
+        (expanded && !block.finished()).then(|| ShellCommandContentExtent {
+            rendered_height: block_content_rows(block).len(),
+            command_rows: usize::from(!block.should_hide_command_grid())
+                .saturating_mul(block.prompt_and_command_grid().len_displayed()),
+            output_rows: usize::from(!block.should_hide_output_grid())
+                .saturating_mul(block.output_grid().len_displayed()),
+        })
+    }
+    /// Whether running command output has changed height since the last layout.
+    pub(super) fn needs_height_measurement(&self) -> bool {
+        self.dynamic_content_extent()
+            .is_some_and(|extent| self.last_measured_content_extent.get() != Some(extent))
+    }
+
+    /// Records the row extent represented by the latest agent-block layout.
+    pub(super) fn record_content_extent_measurement(
+        &self,
+        extent: Option<ShellCommandContentExtent>,
+    ) {
+        self.last_measured_content_extent.set(extent);
     }
     pub(super) fn is_expanded(&self) -> bool {
         !self.state.is_collapsed()
@@ -437,11 +465,9 @@ impl TuiView for TuiShellCommandView {
             .as_ref(app)
             .get_action_status(&self.action.id);
         if matches!(status, Some(AIActionStatus::Blocked)) {
-            self.command_running.set(false);
             return self.render_blocked(app);
         }
         let Some(block) = self.resolved_block(status.as_ref()) else {
-            self.command_running.set(false);
             return render_fallback_tool_call_section(
                 &self.action,
                 status.as_ref(),
@@ -450,8 +476,6 @@ impl TuiView for TuiShellCommandView {
                 app,
             );
         };
-        self.command_running
-            .set(matches!(block.details.state, CommandBlockState::Running));
 
         let builder = TuiUiBuilder::from_app(app);
         let display_state =

crates/warp_tui/src/tui_shell_command_view_tests.rs

@@ -371,6 +371,50 @@ fn terminal_block_is_collapsed_by_default_and_expands_inline() {
     });
 }
 
+#[test]
+fn running_expanded_command_remeasures_only_when_row_extent_changes() {
+    App::test((), |mut app| async move {
+        let action = command_action("action-1", "printf rows");
+        let terminal_model = terminal_model_with_running_command(&action, "printf rows", "partial");
+        let view = add_shell_view(&mut app, action, terminal_model.clone());
+        view.update(&mut app, |view, ctx| {
+            view.handle_action(&TuiShellCommandViewAction::ToggleExpanded, ctx);
+        });
+
+        let initial_extent = app.read(|app| {
+            let view = view.as_ref(app);
+            let extent = view
+                .dynamic_content_extent()
+                .expect("expanded running command has dynamic content");
+            assert!(view.needs_height_measurement());
+            view.record_content_extent_measurement(Some(extent));
+            assert!(!view.needs_height_measurement());
+            extent
+        });
+
+        terminal_model.lock().process_bytes("\rshort");
+        app.read(|app| {
+            assert!(!view.as_ref(app).needs_height_measurement());
+        });
+
+        let added_rows = "next row\r\n".repeat(100);
+        terminal_model.lock().process_bytes(added_rows.as_str());
+        app.read(|app| {
+            let view = view.as_ref(app);
+            let current_extent = view
+                .dynamic_content_extent()
+                .expect("command remains active after receiving output");
+            assert_ne!(current_extent, initial_extent);
+            assert!(view.needs_height_measurement());
+        });
+
+        terminal_model.lock().finish_block();
+        app.read(|app| {
+            assert!(!view.as_ref(app).needs_height_measurement());
+        });
+    });
+}
+
 #[test]
 fn long_path_command_wraps_in_full_with_the_chevron_on_the_first_row() {
     App::test((), |mut app| async move {
@@ -635,6 +679,24 @@ fn terminal_model_with_command(
     Arc::new(FairMutex::new(model))
 }
 
+fn terminal_model_with_running_command(
+    action: &AIAgentAction,
+    command: &str,
+    output: &str,
+) -> Arc<FairMutex<TerminalModel>> {
+    let mut model = TerminalModel::mock(None, None);
+    model.simulate_long_running_block(command, output);
+    model
+        .block_list_mut()
+        .active_block_mut()
+        .set_agent_interaction_mode_for_requested_command(
+            action.id.clone(),
+            None,
+            AIConversationId::new(),
+        );
+    Arc::new(FairMutex::new(model))
+}
+
 fn command_action(id: &str, command: &str) -> AIAgentAction {
     AIAgentAction {
         id: AIAgentActionId::from(id.to_owned()),