Load deferred tool outputs in one request without scroll animation

AGENTS.md

@@ -81,7 +81,7 @@ Layer 5 (binary):     code_assistant    — CLI, config, feature-gated frontend
   in content-addressed `blobs/`; checkpoints append deltas only
 - Resident sessions keep those results unread; snapshots defer their
   outputs (`ToolResultData::output_deferred`), frontends fetch them with
-  `SessionService::load_tool_output`; agent runs resolve all of them
+  `SessionService::load_tool_outputs`; agent runs resolve all of them
 - Flat legacy session files are migrated at startup; see
   `docs/session-storage.md`
 

crates/code_assistant_core/src/session/instance.rs

@@ -168,6 +168,10 @@ pub struct ToolOutputLoader {
 }
 
 impl ToolOutputLoader {
+    pub fn tool_id(&self) -> &str {
+        &self.execution.tool_request.id
+    }
+
     pub fn load(mut self) -> Result<crate::ui::ui_events::ToolResultData> {
         self.persistence
             .resolve_tool_results(&self.session_id, std::slice::from_mut(&mut self.execution))?;
@@ -435,6 +439,11 @@ impl SessionInstance {
 
     /// Terminate the running agent and release the cross-process agent lock.
     pub fn terminate_agent(&mut self) {
+        // Release before stopping: the stopped run resolves its outcome on
+        // another thread, and whoever observes that may rely on the locks
+        // being gone.
+        self.agent_lock = None;
+        self.sleep_guard = None;
         self.request_stop();
         if let Some(handle) = self.setup_task.take() {
             handle.abort();
@@ -443,9 +452,6 @@ impl SessionInstance {
             handle.abort();
             self.clear_fragment_buffer();
         }
-        // Release the cross-process agent lock
-        self.agent_lock = None;
-        self.sleep_guard = None;
         self.set_activity_state(SessionActivityState::Idle);
     }
 

crates/code_assistant_core/src/session/manager.rs

@@ -1592,16 +1592,25 @@ impl SessionManager {
 
     /// Prepare reading the complete UI data of a tool result of a resident
     /// session, to be done without the manager.
-    pub fn tool_output_loader(
+    pub fn tool_output_loaders(
         &self,
         session_id: &str,
-        tool_id: &str,
-    ) -> Result<crate::session::instance::ToolOutputLoader> {
-        self.active_sessions
+        tool_ids: &[String],
+    ) -> Result<Vec<crate::session::instance::ToolOutputLoader>> {
+        let session = self
+            .active_sessions
             .get(session_id)
-            .ok_or_else(|| anyhow::anyhow!("Session not found: {session_id}"))?
-            .tool_output_loader(&self.persistence, tool_id)
-            .ok_or_else(|| anyhow::anyhow!("No tool execution {tool_id} in {session_id}"))
+            .ok_or_else(|| anyhow::anyhow!("Session not found: {session_id}"))?;
+        Ok(tool_ids
+            .iter()
+            .filter_map(|tool_id| {
+                let loader = session.tool_output_loader(&self.persistence, tool_id);
+                if loader.is_none() {
+                    warn!("No tool execution {tool_id} in {session_id}");
+                }
+                loader
+            })
+            .collect())
     }
 
     /// Get a mutable session instance by ID

crates/code_assistant_core/src/session/service.rs

@@ -545,23 +545,39 @@ impl SessionService {
         .await
     }
 
-    /// The complete UI data of a tool result of a loaded session, including
-    /// an output its snapshot deferred
-    /// ([`crate::ui::ui_events::ToolResultData::output_deferred`]). Runs
-    /// apart from the session's commands; the result is read without holding
-    /// up other sessions.
-    pub async fn load_tool_output(
+    /// The complete UI data of tool results of a loaded session, including
+    /// the outputs its snapshot deferred
+    /// ([`crate::ui::ui_events::ToolResultData::output_deferred`]). One call
+    /// for all of them, so a frontend gets them together; a result that
+    /// can't be read is left out. Runs apart from the session's commands;
+    /// the results are read without holding up other sessions.
+    pub async fn load_tool_outputs(
         &self,
         session_id: String,
-        tool_id: String,
-    ) -> Result<crate::ui::ui_events::ToolResultData> {
+        tool_ids: Vec<String>,
+    ) -> Result<Vec<ToolResultData>> {
         self.call_io(move |ctx| async move {
-            let loader = ctx
+            let loaders = ctx
                 .manager
                 .lock()
                 .await
-                .tool_output_loader(&session_id, &tool_id)?;
-            tokio::task::spawn_blocking(move || loader.load()).await?
+                .tool_output_loaders(&session_id, &tool_ids)?;
+            let results = tokio::task::spawn_blocking(move || {
+                loaders
+                    .into_iter()
+                    .filter_map(|loader| {
+                        let tool_id = loader.tool_id().to_string();
+                        loader
+                            .load()
+                            .inspect_err(|e| {
+                                warn!("Failed to load the output of tool {tool_id}: {e:#}")
+                            })
+                            .ok()
+                    })
+                    .collect()
+            })
+            .await?;
+            Ok(results)
         })
         .await
     }

crates/code_assistant_core/src/session/service/stored_output_tests.rs

@@ -125,12 +125,16 @@ async fn showing_a_session_leaves_large_successful_outputs_unread() {
 
     std::fs::write(&blob, blob_content).unwrap();
     let loaded = service
-        .load_tool_output(id.clone(), "large".into())
+        .load_tool_outputs(id.clone(), vec!["large".into(), "unknown".into()])
         .await
         .unwrap();
+    // One call for all of them; a tool the session doesn't know is left out.
+    assert_eq!(loaded.len(), 1);
+    let loaded = &loaded[0];
+    assert_eq!(loaded.tool_id, "large");
     assert!(!loaded.output_deferred);
     assert_eq!(loaded.status, ToolStatus::Success);
-    assert!(loaded.output.unwrap().contains(&large));
+    assert!(loaded.output.as_ref().unwrap().contains(&large));
 }
 
 #[cfg(unix)]

crates/code_assistant_core/src/ui/ui_events.rs

@@ -41,7 +41,7 @@ pub struct ToolResultData {
     pub images: Vec<ImageData>,
     /// The result is complete except for `output`, `styled_output` and
     /// `images`, which are left out because reading it is expensive (it is
-    /// stored outside the session record). `SessionService::load_tool_output`
+    /// stored outside the session record). `SessionService::load_tool_outputs`
     /// returns it complete.
     pub output_deferred: bool,
 }
@@ -243,6 +243,14 @@ pub enum UiEvent {
     ClearMessages,
     /// Update metadata for a single session without refreshing the entire list
     UpdateSessionMetadata { metadata: ChatMetadata },
+    /// Outputs a frontend asked for with `SessionService::load_tool_outputs`
+    /// arrived: complete results for tool blocks shown with
+    /// [`ToolResultData::output_deferred`]. Not new content, so no reason to
+    /// scroll; tool IDs are only unique per session.
+    ToolOutputsLoaded {
+        session_id: String,
+        results: Vec<ToolResultData>,
+    },
     /// A session's lifecycle record changed (visited, settled, un-settled).
     /// Relevant for every view's sidebar, not only the session's own.
     UpdateSessionLifecycle {

crates/ui_acp/src/ui.rs

@@ -985,6 +985,9 @@ impl UserInterface for ACPUserUI {
             UiEvent::UpdateReviewFiles | UiEvent::UpdateReviewDiff => {
                 // Review panel is GPUI-specific.
             }
+            UiEvent::ToolOutputsLoaded { .. } => {
+                // Deferred tool outputs are only asked for by GPUI.
+            }
             UiEvent::UpdateAllowedModels { .. } => {
                 // Model dropdown filtering is GPUI-specific.
             }

crates/ui_gpui/src/app/commands.rs

@@ -160,25 +160,23 @@ impl Gpui {
         });
     }
 
-    /// Load the output a restored tool result left out
-    /// (`ToolResultData::output_deferred`) and apply it like a status update.
-    pub(crate) fn cmd_load_tool_output(&self, session_id: String, tool_id: String) {
+    /// Load the outputs restored tool results left out
+    /// (`ToolResultData::output_deferred`), all in one go.
+    pub(crate) fn cmd_load_tool_outputs(&self, session_id: String, tool_ids: Vec<String>) {
         let Some(service) = self.session_service() else {
             return;
         };
         let gpui = self.clone();
         self.dispatch(async move {
-            match service.load_tool_output(session_id, tool_id.clone()).await {
-                Ok(result) => gpui.push_event(UiEvent::UpdateToolStatus {
-                    tool_id: result.tool_id,
-                    status: result.status,
-                    message: result.message,
-                    output: result.output,
-                    styled_output: result.styled_output,
-                    duration_seconds: result.duration_seconds,
-                    images: result.images,
+            match service
+                .load_tool_outputs(session_id.clone(), tool_ids)
+                .await
+            {
+                Ok(results) => gpui.push_event(UiEvent::ToolOutputsLoaded {
+                    session_id,
+                    results,
                 }),
-                Err(e) => warn!("Failed to load the output of tool {tool_id}: {e:#}"),
+                Err(e) => warn!("Failed to load tool outputs of {session_id}: {e:#}"),
             }
         });
     }

crates/ui_gpui/src/app/event_loop.rs

@@ -75,6 +75,22 @@ impl Gpui {
         self.update_all_messages(cx, |message_container, cx| {
             message_container.apply_tool_results(|tool_id| by_id.get(tool_id).copied(), cx);
         });
+        self.load_shown_deferred_outputs(cx);
+    }
+
+    /// Ask for the outputs that shown tool blocks lack, in one request.
+    fn load_shown_deferred_outputs(&self, cx: &mut gpui_kit::AsyncApp) {
+        let Some(session_id) = self.current_session_id.lock().unwrap().clone() else {
+            return;
+        };
+        let containers = self.message_queue.lock().unwrap().clone();
+        let tool_ids: Vec<String> = containers
+            .iter()
+            .flat_map(|container| cx.update_entity(container, |c, cx| c.claim_deferred_outputs(cx)))
+            .collect();
+        if !tool_ids.is_empty() {
+            self.cmd_load_tool_outputs(session_id, tool_ids);
+        }
     }
 
     pub(crate) fn process_ui_event_async(&self, event: UiEvent, cx: &mut gpui_kit::AsyncApp) {
@@ -230,6 +246,23 @@ impl Gpui {
                 self.auto_scroll_if_following(cx);
             }
 
+            UiEvent::ToolOutputsLoaded {
+                session_id,
+                results,
+            } => {
+                // Asked for by a session no longer shown, whose tool IDs may
+                // recur in this one.
+                if !self.is_current_session(&session_id) {
+                    debug!("Dropping tool outputs loaded for {session_id}");
+                    return;
+                }
+                self.apply_tool_results(&results, cx);
+                self.update_messages_view(cx, |view, cx| {
+                    view.keep_tail_in_view();
+                    cx.notify();
+                });
+            }
+
             UiEvent::EndTool { id } => {
                 self.update_all_messages(cx, |message_container, cx| {
                     message_container.end_tool_use(&id, cx);

crates/ui_gpui/src/blocks/container.rs

@@ -213,6 +213,16 @@ impl MessageContainer {
         elements.clone()
     }
 
+    /// The IDs of the tool blocks that show an output their restored result
+    /// left out and haven't asked for it yet
+    /// ([`BlockView::claim_deferred_output`]).
+    pub fn claim_deferred_outputs(&self, cx: &mut Context<Self>) -> Vec<String> {
+        self.elements()
+            .iter()
+            .filter_map(|view| view.update(cx, |view, _| view.claim_deferred_output()))
+            .collect()
+    }
+
     /// Returns true if this container has no block elements.
     pub fn is_empty(&self) -> bool {
         self.elements.lock().unwrap().is_empty()
@@ -467,7 +477,6 @@ impl MessageContainer {
                 } else if !view.is_generating {
                     view.set_generating(true);
                 }
-                view.load_deferred_output(cx);
 
                 updated = true;
                 cx.notify();

crates/ui_gpui/src/blocks/mod.rs

@@ -227,7 +227,7 @@ pub struct BlockView {
     /// Focus target so a drag-select in the diff card can focus the block and
     /// Cmd/Ctrl-C reaches [`Self::copy_diff_selection`].
     focus_handle: FocusHandle,
-    /// Whether [`Self::load_deferred_output`] asked for the tool's output.
+    /// Whether [`Self::claim_deferred_output`] handed out the tool's ID.
     deferred_output_requested: bool,
 }
 
@@ -286,27 +286,22 @@ impl BlockView {
         }
     }
 
-    /// Ask for the output a restored tool result left out
-    /// ([`ToolUseBlock::output_deferred`]) once the block shows it: when it
-    /// is expanded, which cards are unless collapsed by hand. It arrives as a
-    /// tool status update.
-    pub(crate) fn load_deferred_output(&mut self, cx: &mut Context<Self>) {
-        let Some(tool) = self.block.as_tool() else {
-            return;
-        };
+    /// The tool's ID, once, if the block shows an output a restored tool
+    /// result left out ([`ToolUseBlock::output_deferred`]): when it is
+    /// expanded, which cards are unless collapsed by hand. The caller asks
+    /// for it; it arrives as [`UiEvent::ToolOutputsLoaded`].
+    ///
+    /// [`UiEvent::ToolOutputsLoaded`]: code_assistant_core::ui::UiEvent::ToolOutputsLoaded
+    pub(crate) fn claim_deferred_output(&mut self) -> Option<String> {
+        let tool = self.block.as_tool()?;
         if !tool.output_deferred
             || tool.state != ToolBlockState::Expanded
             || self.deferred_output_requested
         {
-            return;
+            return None;
         }
-        let (Some(session_id), Some(gpui)) =
-            (self.session_id.clone(), cx.try_global::<crate::Gpui>())
-        else {
-            return;
-        };
         self.deferred_output_requested = true;
-        gpui.cmd_load_tool_output(session_id, tool.id.clone());
+        Some(tool.id.clone())
     }
 
     /// The block if it is the tool block `tool_id`, for assertions.
@@ -604,7 +599,12 @@ impl BlockView {
         if let (Some(session_id), Some(tool)) = (&self.session_id, self.block.as_tool()) {
             ToolCollapseState::set(session_id, &tool.id, tool.state.clone(), cx);
         }
-        self.load_deferred_output(cx);
+        if let Some(tool_id) = self.claim_deferred_output()
+            && let (Some(session_id), Some(gpui)) =
+                (self.session_id.clone(), cx.try_global::<crate::Gpui>())
+        {
+            gpui.cmd_load_tool_outputs(session_id, vec![tool_id]);
+        }
 
         self.start_expand_collapse_animation(should_expand, cx);
     }

crates/ui_gpui/src/main_screen/tests/end_to_end.rs

@@ -42,7 +42,11 @@ fn a_message_sent_while_the_agent_works_waits_for_it(cx: &mut TestAppContext) {
     let mut test = MainScreenTest::with_core(cx, waiting_for_the_model());
     test.open_new_session();
     test.send("first");
-    test.wait_until("the agent works", |test| test.agent_is_running());
+    // Not just running: until its request is sent, the agent takes a waiting
+    // message into it.
+    test.wait_until("the model is asked", |test| {
+        test.agent_waits_for_the_model()
+    });
 
     test.send("second");
     test.wait_until("the message waits", |test| test.pending_message().is_some());