fix: wait for agent prompt readiness (#2537)

* fix: wait for agent prompt readiness

refs #2410

* fix: return pane updates after tab bar command

refs #2410

---------

Co-authored-by: Can Celik <ogulcancelik@gmail.com>

2 files changed, 151 insertions, 18 deletions

diff --git a/src/app/actions.rs b/src/app/actions.rs
index 5d8051fd..d6e264f5 100644
--- a/src/app/actions.rs
+++ b/src/app/actions.rs
@@ -177,7 +177,10 @@ pub fn notification_toast_for_pane_state_update(
     suppress_active_tab_notifications: bool,
     update: &PaneStateUpdate,
 ) -> Option<ToastKind> {
-    if suppress_active_tab_notifications || update.state == update.previous_state {
+    if update.suppress_completion
+        || suppress_active_tab_notifications
+        || update.state == update.previous_state
+    {
         return None;
     }
 
@@ -249,6 +252,7 @@ pub struct PaneStateUpdate {
     pub agent_name_changed: bool,
     pub agent_released: bool,
     pub agent_release_status: Option<crate::api::schema::AgentStatus>,
+    pub suppress_completion: bool,
 }
 
 // ---------------------------------------------------------------------------
@@ -1036,7 +1040,7 @@ impl AppState {
                     .get_mut(&terminal_id)?
                     .expire_agent_metadata_at(scheduled_deadline, now)?;
                 let change = mutation.effective_state_change?;
-                let seen = self.apply_pane_state_change(ws_idx, pane_id, &change)?;
+                let seen = self.apply_pane_state_change(ws_idx, pane_id, &change, false)?;
                 let update = PaneStateUpdate {
                     pane_id,
                     ws_idx,
@@ -1053,6 +1057,7 @@ impl AppState {
                     agent_name_changed: false,
                     agent_released: false,
                     agent_release_status: None,
+                    suppress_completion: false,
                 };
                 Some(update)
             })
@@ -2771,6 +2776,16 @@ impl AppState {
                 }
                 Vec::new()
             }
+            AppEvent::AgentProcessDetected {
+                pane_id,
+                agent,
+                observed_at,
+            } => self
+                .update_terminal_state(pane_id, |terminal| {
+                    Some(terminal.set_detected_agent_process_at(agent, observed_at))
+                })
+                .into_iter()
+                .collect(),
             AppEvent::StateChanged {
                 pane_id,
                 agent,
@@ -2954,11 +2969,20 @@ impl AppState {
             .clone();
         let previous_seen = self.workspaces[ws_idx].pane_state(pane_id)?.seen;
         let now = Instant::now();
-        let (mutation, managed_changed, agent_name_changed, unchanged_change) = {
+        let (
+            mutation,
+            managed_changed,
+            agent_name_changed,
+            unchanged_change,
+            managed_launch_pending,
+            suppress_acquisition_completion,
+        ) = {
             let terminal = self.terminals.get_mut(&terminal_id)?;
             let previous_agent_name = terminal.agent_name.clone();
+            let managed_launch_pending = terminal.managed_agent_launch_pending();
             let mutation = update(terminal)?;
             let managed_changed = terminal.reconcile_managed_agent_at(now, false);
+            let suppress_acquisition_completion = terminal.finish_agent_process_acquisition();
             let agent_name_changed = terminal.agent_name != previous_agent_name;
             let unchanged_change = (mutation.agent_released || agent_name_changed)
                 .then(|| terminal.unchanged_effective_state_change_at(now));
@@ -2967,6 +2991,8 @@ impl AppState {
                 managed_changed,
                 agent_name_changed,
                 unchanged_change,
+                managed_launch_pending,
+                suppress_acquisition_completion,
             )
         };
         if mutation.session_ref_changed || managed_changed || agent_name_changed {
@@ -2974,13 +3000,15 @@ impl AppState {
         }
         let agent_released = mutation.agent_released;
         let change = mutation.effective_state_change.or(unchanged_change)?;
+        let suppress_completion = change.state == AgentState::Idle
+            && (managed_launch_pending || suppress_acquisition_completion);
         if change.previous_state != change.state {
             self.next_agent_state_change_seq += 1;
             if let Some(terminal) = self.terminals.get_mut(&terminal_id) {
                 terminal.last_agent_state_change_seq = Some(self.next_agent_state_change_seq);
             }
         }
-        let seen = self.apply_pane_state_change(ws_idx, pane_id, &change)?;
+        let seen = self.apply_pane_state_change(ws_idx, pane_id, &change, suppress_completion)?;
         let update = PaneStateUpdate {
             pane_id,
             ws_idx,
@@ -3005,6 +3033,7 @@ impl AppState {
             agent_name_changed,
             agent_released,
             agent_release_status: agent_released.then(|| pane_agent_status(change.state, seen)),
+            suppress_completion,
         };
         Some(update)
     }
@@ -3071,6 +3100,7 @@ impl AppState {
         ws_idx: usize,
         pane_id: PaneId,
         change: &EffectiveStateChange,
+        suppress_completion: bool,
     ) -> Option<bool> {
         let is_active_tab = self.pane_is_in_active_tab(ws_idx, pane_id);
         let suppress_active_tab_notifications =
@@ -3082,13 +3112,17 @@ impl AppState {
 
         if change.state != AgentState::Idle {
             pane.seen = true;
-        } else if is_completion_transition(change) {
+        } else if !suppress_completion && is_completion_transition(change) {
             pane.seen = suppress_active_tab_notifications;
         }
         let seen = pane.seen;
 
-        if let Some(delivery) = self.record_or_deliver_agent_notification(ws_idx, pane_id, change) {
-            self.apply_agent_notification_delivery(&delivery);
+        if !suppress_completion {
+            if let Some(delivery) =
+                self.record_or_deliver_agent_notification(ws_idx, pane_id, change)
+            {
+                self.apply_agent_notification_delivery(&delivery);
+            }
         }
 
         Some(seen)
@@ -4965,6 +4999,97 @@ mod tests {
     }
 
     #[test]
+    fn first_idle_after_process_detection_is_not_completion() {
+        let mut state = app_with_workspaces(&["active", "background"]);
+        state.toast_config.delivery = crate::config::ToastDelivery::Herdr;
+        state.active = Some(0);
+        let pane_id = *state.workspaces[1].panes.keys().next().unwrap();
+
+        state.handle_app_event(AppEvent::AgentProcessDetected {
+            pane_id,
+            agent: Agent::Pi,
+            observed_at: Instant::now(),
+        });
+        let direct_idle = state
+            .handle_app_event(AppEvent::StateChanged {
+                pane_id,
+                agent: Some(Agent::Pi),
+                state: AgentState::Idle,
+                visible_blocker: false,
+                visible_working: false,
+                process_exited: false,
+                observed_at: Instant::now(),
+            })
+            .pop()
+            .expect("direct idle state update");
+        assert!(direct_idle.suppress_completion);
+
+        state.handle_app_event(AppEvent::AgentProcessDetected {
+            pane_id,
+            agent: Agent::Pi,
+            observed_at: Instant::now(),
+        });
+        for agent_state in [AgentState::Working, AgentState::Blocked] {
+            state.handle_app_event(AppEvent::StateChanged {
+                pane_id,
+                agent: Some(Agent::Pi),
+                state: agent_state,
+                visible_blocker: agent_state == AgentState::Blocked,
+                visible_working: agent_state == AgentState::Working,
+                process_exited: false,
+                observed_at: Instant::now(),
+            });
+        }
+        let update = state
+            .handle_app_event(AppEvent::StateChanged {
+                pane_id,
+                agent: Some(Agent::Pi),
+                state: AgentState::Idle,
+                visible_blocker: false,
+                visible_working: false,
+                process_exited: false,
+                observed_at: Instant::now(),
+            })
+            .pop()
+            .expect("idle state update");
+
+        assert!(update.suppress_completion);
+        assert!(state.workspaces[1].panes[&pane_id].seen);
+        assert!(!matches!(
+            state.toast.as_ref().map(|toast| toast.kind),
+            Some(ToastKind::Finished)
+        ));
+
+        state.handle_app_event(AppEvent::AgentProcessDetected {
+            pane_id,
+            agent: Agent::Codex,
+            observed_at: Instant::now(),
+        });
+        state.handle_app_event(AppEvent::StateChanged {
+            pane_id,
+            agent: Some(Agent::Codex),
+            state: AgentState::Working,
+            visible_blocker: false,
+            visible_working: true,
+            process_exited: false,
+            observed_at: Instant::now(),
+        });
+        let exit_update = state
+            .handle_app_event(AppEvent::StateChanged {
+                pane_id,
+                agent: Some(Agent::Codex),
+                state: AgentState::Idle,
+                visible_blocker: false,
+                visible_working: false,
+                process_exited: true,
+                observed_at: Instant::now(),
+            })
+            .pop()
+            .expect("process exit update");
+        assert!(!exit_update.suppress_completion);
+    }
+
+    #[test]
     fn waiting_sound_plays_even_in_active_workspace() {
         assert_eq!(
             notification_sound_for_state_change(true, AgentState::Working, AgentState::Blocked),
diff --git a/src/app/api.rs b/src/app/api.rs
index 18471f57..92d9596f 100644
--- a/src/app/api.rs
+++ b/src/app/api.rs
@@ -106,13 +106,20 @@ impl App {
     }
 
     pub(crate) fn handle_internal_event(&mut self, ev: AppEvent) {
+        let _ = self.handle_internal_event_with_pane_updates(ev);
+    }
+
+    pub(crate) fn handle_internal_event_with_pane_updates(
+        &mut self,
+        ev: AppEvent,
+    ) -> Vec<crate::app::actions::PaneStateUpdate> {
         if let AppEvent::TerminalBell { count, .. } = ev {
             if let Err(err) =
                 crate::terminal_effects::write_terminal_bells(&mut std::io::stdout(), count)
             {
                 tracing::warn!(err = %err, "failed to emit terminal bell");
             }
-            return;
+            return Vec::new();
         }
 
         if let AppEvent::ClipboardWrite { content } = ev {
@@ -121,7 +128,7 @@ impl App {
             #[cfg(test)]
             let _ = content;
             self.show_clipboard_feedback();
-            return;
+            return Vec::new();
         }
 
         if let AppEvent::PrefixInputSource { active } = ev {
@@ -130,14 +137,14 @@ impl App {
             // App-internal drain consume the event before the forwarding drain, the flag keeps the
             // switch out of the headless server process.
             if !self.local_input_source_switch {
-                return;
+                return Vec::new();
             }
             if active {
                 self.prefix_input_source.switch_to_ascii();
             } else {
                 self.prefix_input_source.restore();
             }
-            return;
+            return Vec::new();
         }
 
         if let AppEvent::GitStatusRefreshed {
@@ -146,7 +153,7 @@ impl App {
         } = ev
         {
             self.handle_git_status_refreshed(results, cache_updates);
-            return;
+            return Vec::new();
         }
 
         if let AppEvent::TabBarCommandFinished {
@@ -156,7 +163,7 @@ impl App {
         } = ev
         {
             let _ = self.handle_tab_bar_command_finished(generation, segment_index, result);
-            return;
+            return Vec::new();
         }
 
         if let AppEvent::PluginCommandFinished {
@@ -187,17 +194,17 @@ impl App {
                     crate::api::schema::PluginCommandStatus::Failed
                 };
             }
-            return;
+            return Vec::new();
         }
 
         if let AppEvent::WorktreeAddFinished(result) = ev {
             self.handle_worktree_add_finished(*result);
-            return;
+            return Vec::new();
         }
 
         if let AppEvent::WorktreeRemoveFinished(result) = ev {
             self.handle_worktree_remove_finished(*result);
-            return;
+            return Vec::new();
         }
 
         if let AppEvent::PaneDied { pane_id } = &ev {
@@ -208,7 +215,7 @@ impl App {
                 .is_some_and(|popup| popup.pane_id == *pane_id)
             {
                 self.close_popup_pane();
-                return;
+                return Vec::new();
             }
             let previous_toast = self.state.toast.clone();
             if let Some(update) = self.state.publish_pane_process_exit_if_agent(*pane_id) {
@@ -223,7 +230,7 @@ impl App {
                 self.overlay_panes.remove(pane_id);
                 self.render_dirty.request_generic();
                 self.render_notify.notify_one();
-                return;
+                return Vec::new();
             }
         }
 
@@ -368,6 +375,7 @@ impl App {
 
         self.sync_toast_deadline(previous_toast);
         self.shutdown_detached_terminal_runtimes();
+        pane_updates
     }
 
     fn reset_agent_detection_for_agents(&self, agents: &[crate::detect::Agent]) {