Render scheduled wakeups as an inline tool line

crates/code_assistant_core/src/mocks.rs

@@ -1223,12 +1223,13 @@ impl ToolTestFixture {
         struct NullSink;
         #[async_trait::async_trait]
         impl crate::session::wakeup::WakeupSink for NullSink {
-            async fn fire(&self, _session_id: &str, _message: &str) {}
+            async fn fire(&self, _session_id: &str, _wakeup_id: u64, _message: &str) {}
         }
         let handle = crate::session::wakeup::spawn_wakeup_scheduler(NullSink, None);
         self.services.wakeups = Some(crate::session::wakeup::SessionWakeups {
             handle,
             session_id: "test-session".to_string(),
+            used_ids: 0,
         });
         self
     }

crates/code_assistant_core/src/session/instance.rs

@@ -668,6 +668,8 @@ impl SessionInstance {
             pending_permission_requests: self.pending_permission_requests.snapshot(),
             pending_new_context_target: self.pending_new_context_target.snapshot(),
             pending_questions: self.pending_questions.snapshot(),
+            // Filled in by the SessionManager, which holds the scheduler.
+            armed_wakeups: Vec::new(),
         })
     }
 

crates/code_assistant_core/src/session/lifecycle.rs

@@ -190,6 +190,9 @@ pub enum SessionStatus {
     Working,
     /// The agent waits out a rate limit.
     RateLimited,
+    /// The agent stopped, but a wakeup it scheduled will continue the
+    /// session; nothing to do for the user.
+    Scheduled,
     /// The agent stopped with an error.
     Failed,
     /// Another code-assistant process runs this session.
@@ -200,13 +203,19 @@ pub enum SessionStatus {
 }
 
 impl SessionStatus {
-    pub fn resolve(activity: &SessionActivityState, awaiting: Option<AwaitingUser>) -> Self {
+    /// `wakeup_armed`: a wakeup is still to fire in the session.
+    pub fn resolve(
+        activity: &SessionActivityState,
+        awaiting: Option<AwaitingUser>,
+        wakeup_armed: bool,
+    ) -> Self {
         match awaiting {
             Some(AwaitingUser::Approval) => return Self::NeedsApproval,
             Some(AwaitingUser::Answer) => return Self::NeedsAnswer,
             None => {}
         }
         match activity {
+            SessionActivityState::Idle if wakeup_armed => Self::Scheduled,
             SessionActivityState::Idle => Self::Ready,
             SessionActivityState::AgentRunning | SessionActivityState::WaitingForResponse => {
                 Self::Working
@@ -225,7 +234,7 @@ impl SessionStatus {
             return false;
         }
         match self {
-            Self::Working | Self::RateLimited | Self::RunningElsewhere => true,
+            Self::Working | Self::RateLimited | Self::Scheduled | Self::RunningElsewhere => true,
             Self::Ready => !unread,
             Self::NeedsApproval | Self::NeedsAnswer | Self::Failed => false,
         }
@@ -458,24 +467,40 @@ mod tests {
             },
         ] {
             assert_eq!(
-                SessionStatus::resolve(&activity, Some(AwaitingUser::Approval)),
+                SessionStatus::resolve(&activity, Some(AwaitingUser::Approval), true),
                 SessionStatus::NeedsApproval
             );
             assert_eq!(
-                SessionStatus::resolve(&activity, Some(AwaitingUser::Answer)),
+                SessionStatus::resolve(&activity, Some(AwaitingUser::Answer), true),
                 SessionStatus::NeedsAnswer
             );
         }
         assert_eq!(
-            SessionStatus::resolve(&SessionActivityState::WaitingForResponse, None),
+            SessionStatus::resolve(&SessionActivityState::WaitingForResponse, None, true),
             SessionStatus::Working
         );
         assert_eq!(
-            SessionStatus::resolve(&SessionActivityState::Idle, None),
+            SessionStatus::resolve(&SessionActivityState::Idle, None, false),
             SessionStatus::Ready
         );
     }
 
+    #[test]
+    fn idle_session_with_a_wakeup_to_fire_is_scheduled() {
+        assert_eq!(
+            SessionStatus::resolve(&SessionActivityState::Idle, None, true),
+            SessionStatus::Scheduled
+        );
+        let failed = SessionActivityState::Errored {
+            message: "x".into(),
+        };
+        assert_eq!(
+            SessionStatus::resolve(&failed, None, true),
+            SessionStatus::Failed
+        );
+        assert!(SessionStatus::Scheduled.should_recede(true, false));
+    }
+
     #[test]
     fn busy_rows_recede_and_read_ready_rows_recede() {
         assert!(SessionStatus::Working.should_recede(true, false));

crates/code_assistant_core/src/session/manager.rs

@@ -703,6 +703,11 @@ impl SessionManager {
         };
 
         snapshot.current_model = model_name_for_event;
+        snapshot.armed_wakeups = self
+            .wakeup_handle
+            .as_ref()
+            .map(|handle| handle.armed(&session_id))
+            .unwrap_or_default();
         snapshot.allowed_models = self
             .allowed_models_for_session(&session_id)
             .unwrap_or_default();
@@ -1309,6 +1314,15 @@ impl SessionManager {
                 crate::session::wakeup::SessionWakeups {
                     handle,
                     session_id: session_id.to_string(),
+                    used_ids: self
+                        .active_sessions
+                        .get(session_id)
+                        .map(|instance| {
+                            crate::session::wakeup::highest_wakeup_id(
+                                &instance.session.tool_executions,
+                            )
+                        })
+                        .unwrap_or(0),
                 }
             }),
             pty_sessions: self
@@ -1590,6 +1604,60 @@ impl SessionManager {
     }
 
     /// Get a session instance by ID
+    /// Record in the stored result of `schedule_wakeup` that wakeup
+    /// `wakeup_id` fired, and return the result's new UI data. `None` when
+    /// this process armed no such wakeup in the session.
+    pub fn mark_wakeup_fired(
+        &mut self,
+        session_id: &str,
+        wakeup_id: u64,
+    ) -> Result<Option<crate::ui::ui_events::ToolResultData>> {
+        use crate::tools::core::{Render, ResourcesTracker};
+        use crate::tools::impls::wakeup::ScheduleWakeupOutput;
+
+        let mut fired = None;
+        let session = self.persistence.update_entry(session_id, |session| {
+            let armed_here = |output: &ScheduleWakeupOutput| {
+                output.wakeup_id == wakeup_id
+                    && output.armed_by.as_deref()
+                        == Some(crate::session::wakeup::process_instance())
+            };
+            for exec in &mut session.tool_executions {
+                if exec.tool_name != "schedule_wakeup" {
+                    continue;
+                }
+                let Ok(mut output) =
+                    serde_json::from_value::<ScheduleWakeupOutput>(exec.result_json.clone())
+                else {
+                    continue;
+                };
+                if armed_here(&output) {
+                    output.fired_at = Some(chrono::Utc::now());
+                    exec.result_json = serde_json::to_value(&output)?;
+                    fired = Some((exec.tool_request.id.clone(), output));
+                    break;
+                }
+            }
+            Ok(())
+        })?;
+        if let Some(instance) = self.active_sessions.get_mut(session_id) {
+            instance.session = session;
+        }
+        Ok(fired.map(|(tool_id, output)| {
+            let mut tracker = ResourcesTracker::new();
+            crate::ui::ui_events::ToolResultData {
+                tool_id,
+                status: crate::ui::ToolStatus::Success,
+                message: Some(output.status()),
+                output: Some(output.render_for_ui(&mut tracker)),
+                styled_output: None,
+                duration_seconds: None,
+                images: Vec::new(),
+                output_deferred: false,
+            }
+        }))
+    }
+
     pub fn get_session(&self, session_id: &str) -> Option<&SessionInstance> {
         self.active_sessions.get(session_id)
     }
@@ -2129,7 +2197,13 @@ impl SessionManager {
                 .active_sessions
                 .get(&metadata.id)
                 .is_some_and(|instance| !instance.get_activity_state().is_terminal());
-            let busy = running_here || self.is_agent_locked_externally(&metadata.id);
+            // A wakeup still to fire continues the session.
+            let wakeup_armed = self
+                .wakeup_handle
+                .as_ref()
+                .is_some_and(|handle| !handle.armed(&metadata.id).is_empty());
+            let busy =
+                running_here || wakeup_armed || self.is_agent_locked_externally(&metadata.id);
             let repo_root = metadata
                 .branch
                 .as_ref()

crates/code_assistant_core/src/session/mod.rs

@@ -73,6 +73,8 @@ pub struct SessionSnapshot {
     pub pending_new_context_target: Option<new_context::NewContextTargetRequest>,
     /// `ask_question` requests still awaiting answers.
     pub pending_questions: Vec<questions::UserQuestionRequest>,
+    /// Wakeups still to fire, soonest first.
+    pub armed_wakeups: Vec<wakeup::ArmedWakeupInfo>,
 }
 
 impl SessionSnapshot {
@@ -96,6 +98,10 @@ impl SessionSnapshot {
                 session_id: self.session_id.clone(),
                 activity_state: self.activity_state.clone(),
             },
+            UiEvent::UpdateArmedWakeups {
+                session_id: self.session_id.clone(),
+                wakeups: self.armed_wakeups.clone(),
+            },
         ];
         // Show the error banner when connecting to an errored session.
         if let instance::SessionActivityState::Errored { message } = &self.activity_state {