From 0ccdd9bf0942555c4d200519156e0f541a4e5ab1 Mon Sep 17 00:00:00 2001 From: Danilo Leal <67129314+danilo-leal@users.noreply.github.com> Date: Fri, 13 Mar 2026 17:01:00 -0300 Subject: [PATCH] agent_ui: Auto-expand and then collapse thinking blocks (#51525) With these newer models that come with different thinking levels, it's become more frequent to want to see what the thinking is outputting. Thus far in Zed, the thinking block would show up automatically collapsed and every time you wanted to see it, you had to expand it manually. This PR changes that by making the thinking block automatically _expanded_ instead, but as soon as it's done, it collapses again. Release Notes: - Agent: Improved visibility of thinking blocks by making them auto-expanded while in progress. --- crates/agent_ui/src/connection_view.rs | 10 ++- .../src/connection_view/thread_view.rs | 73 +++++++++++++++---- 2 files changed, 64 insertions(+), 19 deletions(-) diff --git a/crates/agent_ui/src/connection_view.rs b/crates/agent_ui/src/connection_view.rs index d2226e675a6a242588074dd2e7b646a7376c8c37..4d352c6a8494f97358ee012740e539c750308886 100644 --- a/crates/agent_ui/src/connection_view.rs +++ b/crates/agent_ui/src/connection_view.rs @@ -1263,13 +1263,14 @@ impl ConnectionView { } } AcpThreadEvent::EntryUpdated(index) => { - if let Some(entry_view_state) = self - .thread_view(&thread_id) - .map(|active| active.read(cx).entry_view_state.clone()) - { + if let Some(active) = self.thread_view(&thread_id) { + let entry_view_state = active.read(cx).entry_view_state.clone(); entry_view_state.update(cx, |view_state, cx| { view_state.sync_entry(*index, thread, window, cx) }); + active.update(cx, |active, cx| { + active.auto_expand_streaming_thought(cx); + }); } } AcpThreadEvent::EntriesRemoved(range) => { @@ -1301,6 +1302,7 @@ impl ConnectionView { if let Some(active) = self.thread_view(&thread_id) { active.update(cx, |active, _cx| { active.thread_retry_status.take(); + active.clear_auto_expand_tracking(); }); } if is_subagent { diff --git a/crates/agent_ui/src/connection_view/thread_view.rs b/crates/agent_ui/src/connection_view/thread_view.rs index 35df60b567de86762a9af330013df0fab35f3f01..eed8de86c841350d507b040287088989ae23c023 100644 --- a/crates/agent_ui/src/connection_view/thread_view.rs +++ b/crates/agent_ui/src/connection_view/thread_view.rs @@ -194,6 +194,7 @@ pub struct ThreadView { pub expanded_tool_calls: HashSet, pub expanded_tool_call_raw_inputs: HashSet, pub expanded_thinking_blocks: HashSet<(usize, usize)>, + auto_expanded_thinking_block: Option<(usize, usize)>, pub subagent_scroll_handles: RefCell>, pub edits_expanded: bool, pub plan_expanded: bool, @@ -425,6 +426,7 @@ impl ThreadView { expanded_tool_calls: HashSet::default(), expanded_tool_call_raw_inputs: HashSet::default(), expanded_thinking_blocks: HashSet::default(), + auto_expanded_thinking_block: None, subagent_scroll_handles: RefCell::new(HashMap::default()), edits_expanded: false, plan_expanded: false, @@ -4573,6 +4575,53 @@ impl ThreadView { .into_any_element() } + /// If the last entry's last chunk is a streaming thought block, auto-expand it. + /// Also collapses the previously auto-expanded block when a new one starts. + pub(crate) fn auto_expand_streaming_thought(&mut self, cx: &mut Context) { + let key = { + let thread = self.thread.read(cx); + if thread.status() != ThreadStatus::Generating { + return; + } + let entries = thread.entries(); + let last_ix = entries.len().saturating_sub(1); + match entries.get(last_ix) { + Some(AgentThreadEntry::AssistantMessage(msg)) => match msg.chunks.last() { + Some(AssistantMessageChunk::Thought { .. }) => { + Some((last_ix, msg.chunks.len() - 1)) + } + _ => None, + }, + _ => None, + } + }; + + if let Some(key) = key { + if self.auto_expanded_thinking_block != Some(key) { + if let Some(old_key) = self.auto_expanded_thinking_block.replace(key) { + self.expanded_thinking_blocks.remove(&old_key); + } + self.expanded_thinking_blocks.insert(key); + cx.notify(); + } + } else if self.auto_expanded_thinking_block.is_some() { + // The last chunk is no longer a thought (model transitioned to responding), + // so collapse the previously auto-expanded block. + self.collapse_auto_expanded_thinking_block(); + cx.notify(); + } + } + + fn collapse_auto_expanded_thinking_block(&mut self) { + if let Some(key) = self.auto_expanded_thinking_block.take() { + self.expanded_thinking_blocks.remove(&key); + } + } + + pub(crate) fn clear_auto_expand_tracking(&mut self) { + self.auto_expanded_thinking_block = None; + } + fn render_thinking_block( &self, entry_ix: usize, @@ -4594,20 +4643,6 @@ impl ThreadView { .entry(entry_ix) .and_then(|entry| entry.scroll_handle_for_assistant_message_chunk(chunk_ix)); - let thinking_content = { - div() - .id(("thinking-content", chunk_ix)) - .when_some(scroll_handle, |this, scroll_handle| { - this.track_scroll(&scroll_handle) - }) - .text_ui_sm(cx) - .overflow_hidden() - .child(self.render_markdown( - chunk, - MarkdownStyle::themed(MarkdownFont::Agent, window, cx), - )) - }; - v_flex() .gap_1() .child( @@ -4663,11 +4698,19 @@ impl ThreadView { .when(is_open, |this| { this.child( div() + .id(("thinking-content", chunk_ix)) .ml_1p5() .pl_3p5() .border_l_1() .border_color(self.tool_card_border_color(cx)) - .child(thinking_content), + .when_some(scroll_handle, |this, scroll_handle| { + this.track_scroll(&scroll_handle) + }) + .overflow_hidden() + .child(self.render_markdown( + chunk, + MarkdownStyle::themed(MarkdownFont::Agent, window, cx), + )), ) }) .into_any_element()