fix: deliver agent status subscriptions from event hub

5 files changed, 319 insertions, 39 deletions

diff --git a/src/api/event_hub.rs b/src/api/event_hub.rs
index b66a2157..42cf47a2 100644
--- a/src/api/event_hub.rs
+++ b/src/api/event_hub.rs
@@ -36,4 +36,11 @@ impl EventHub {
             .cloned()
             .collect()
     }
+
+    pub fn current_sequence(&self) -> u64 {
+        let Ok(state) = self.inner.lock() else {
+            return 0;
+        };
+        state.next_sequence
+    }
 }
diff --git a/src/api/schema.rs b/src/api/schema.rs
index 675d0cbc..6eed4d2c 100644
--- a/src/api/schema.rs
+++ b/src/api/schema.rs
@@ -975,6 +975,8 @@ pub enum EventData {
         workspace_id: String,
         agent_status: AgentStatus,
         #[serde(default, skip_serializing_if = "Option::is_none")]
+        agent: Option<String>,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
         title: Option<String>,
         #[serde(default, skip_serializing_if = "Option::is_none")]
         display_agent: Option<String>,
diff --git a/src/api/subscriptions.rs b/src/api/subscriptions.rs
index aaa7e5ac..9c92a41b 100644
--- a/src/api/subscriptions.rs
+++ b/src/api/subscriptions.rs
@@ -49,8 +49,9 @@ pub(super) struct ActiveAgentStatusChangedSubscription {
     pane_id: String,
     status_filter: Option<crate::api::schema::AgentStatus>,
     last_status: Option<crate::api::schema::AgentStatus>,
-    emit_initial_match: bool,
     last_presentation: Option<PanePresentationSnapshot>,
+    last_sequence: u64,
+    initial_event: Option<PaneAgentStatusChangedEvent>,
     request_prefix: String,
 }
 
@@ -71,6 +72,20 @@ impl PanePresentationSnapshot {
             state_labels: pane.state_labels.clone(),
         }
     }
+
+    fn from_event(
+        title: &Option<String>,
+        display_agent: &Option<String>,
+        custom_status: &Option<String>,
+        state_labels: &std::collections::HashMap<String, String>,
+    ) -> Self {
+        Self {
+            title: title.clone(),
+            display_agent: display_agent.clone(),
+            custom_status: custom_status.clone(),
+            state_labels: state_labels.clone(),
+        }
+    }
 }
 
 pub(super) struct ActiveEventSubscription {
@@ -81,7 +96,7 @@ pub(super) struct ActiveEventSubscription {
 pub(super) enum ActiveSubscription {
     Event(ActiveEventSubscription),
     OutputMatched(ActiveOutputMatchedSubscription),
-    AgentStatusChanged(ActiveAgentStatusChangedSubscription),
+    AgentStatusChanged(Box<ActiveAgentStatusChangedSubscription>),
 }
 
 impl ActiveSubscription {
@@ -90,7 +105,7 @@ impl ActiveSubscription {
         request_id: &str,
         index: usize,
         api_tx: &ApiRequestSender,
-        _event_hub: &EventHub,
+        event_hub: &EventHub,
     ) -> Result<Self, ErrorResponse> {
         match subscription {
             Subscription::WorkspaceCreated {} => Ok(Self::Event(ActiveEventSubscription {
@@ -197,20 +212,34 @@ impl ActiveSubscription {
                 pane_id,
                 agent_status,
             } => {
+                let last_sequence = event_hub.current_sequence();
                 let probe = pane_get(format!("{request_id}:sub:{index}:probe"), &pane_id, api_tx)?;
-                let emit_initial_match =
-                    agent_status.is_some_and(|wanted| wanted == probe.agent_status);
+                let last_status = probe.agent_status;
+                let last_presentation = PanePresentationSnapshot::from(&probe);
+                let initial_event = agent_status
+                    .is_some_and(|wanted| wanted == probe.agent_status)
+                    .then_some(PaneAgentStatusChangedEvent {
+                        pane_id: probe.pane_id,
+                        workspace_id: probe.workspace_id,
+                        agent_status: probe.agent_status,
+                        agent: probe.agent,
+                        title: probe.title,
+                        display_agent: probe.display_agent,
+                        custom_status: probe.custom_status,
+                        state_labels: probe.state_labels,
+                    });
 
-                Ok(Self::AgentStatusChanged(
+                Ok(Self::AgentStatusChanged(Box::new(
                     ActiveAgentStatusChangedSubscription {
                         pane_id,
                         status_filter: agent_status,
-                        last_status: Some(probe.agent_status),
-                        emit_initial_match,
-                        last_presentation: Some(PanePresentationSnapshot::from(&probe)),
+                        last_status: Some(last_status),
+                        last_presentation: Some(last_presentation),
+                        last_sequence,
+                        initial_event,
                         request_prefix: format!("{request_id}:sub:{index}"),
                     },
-                ))
+                )))
             }
         }
     }
@@ -226,7 +255,7 @@ impl ActiveSubscription {
                 serde_json::to_value(subscription.poll(api_tx)?).ok()
             }
             Self::AgentStatusChanged(subscription) => {
-                serde_json::to_value(subscription.poll(api_tx)?).ok()
+                serde_json::to_value(subscription.poll(api_tx, event_hub)?).ok()
             }
         }
     }
@@ -281,13 +310,100 @@ impl ActiveOutputMatchedSubscription {
 }
 
 impl ActiveAgentStatusChangedSubscription {
-    fn poll(&mut self, api_tx: &ApiRequestSender) -> Option<SubscriptionEventEnvelope> {
+    fn poll(
+        &mut self,
+        api_tx: &ApiRequestSender,
+        event_hub: &EventHub,
+    ) -> Option<SubscriptionEventEnvelope> {
+        let mut saw_status_event = false;
+        for (sequence, event) in event_hub.events_after(self.last_sequence) {
+            self.last_sequence = sequence;
+            let crate::api::schema::EventData::PaneAgentStatusChanged {
+                pane_id,
+                workspace_id,
+                agent_status,
+                agent,
+                title,
+                display_agent,
+                custom_status,
+                state_labels,
+            } = event.data
+            else {
+                continue;
+            };
+            if event.event != crate::api::schema::EventKind::PaneAgentStatusChanged {
+                continue;
+            }
+            if pane_id != self.pane_id {
+                continue;
+            }
+            saw_status_event = true;
+
+            let current_presentation = PanePresentationSnapshot::from_event(
+                &title,
+                &display_agent,
+                &custom_status,
+                &state_labels,
+            );
+            self.last_status = Some(agent_status);
+            self.last_presentation = Some(current_presentation);
+            if self
+                .status_filter
+                .is_some_and(|wanted| wanted != agent_status)
+            {
+                continue;
+            }
+
+            self.initial_event = None;
+            return Some(SubscriptionEventEnvelope {
+                event: SubscriptionEventKind::PaneAgentStatusChanged,
+                data: SubscriptionEventData::PaneAgentStatusChanged(PaneAgentStatusChangedEvent {
+                    pane_id,
+                    workspace_id,
+                    agent_status,
+                    agent,
+                    title,
+                    display_agent,
+                    custom_status,
+                    state_labels,
+                }),
+            });
+        }
+
+        if saw_status_event {
+            self.initial_event = None;
+        } else if event_hub.current_sequence() != self.last_sequence {
+            return None;
+        } else if let Some(event) = self.initial_event.take() {
+            return Some(SubscriptionEventEnvelope {
+                event: SubscriptionEventKind::PaneAgentStatusChanged,
+                data: SubscriptionEventData::PaneAgentStatusChanged(event),
+            });
+        }
+
+        let before_snapshot_sequence = self.last_sequence;
         let pane = pane_get(
             format!("{}:pane", self.request_prefix),
             &self.pane_id,
             api_tx,
         )
         .ok()?;
+        let after_snapshot_sequence = event_hub.current_sequence();
+        if after_snapshot_sequence != before_snapshot_sequence {
+            return None;
+        }
+
+        let event = self.event_from_snapshot(pane);
+        if event.is_some() {
+            self.last_sequence = after_snapshot_sequence;
+        }
+        event
+    }
+
+    fn event_from_snapshot(
+        &mut self,
+        pane: crate::api::schema::PaneInfo,
+    ) -> Option<SubscriptionEventEnvelope> {
         let current_status = pane.agent_status;
         let current_presentation = PanePresentationSnapshot::from(&pane);
         let previous_status = self.last_status.replace(current_status);
@@ -296,17 +412,13 @@ impl ActiveAgentStatusChangedSubscription {
             .as_ref()
             .is_some_and(|previous| previous != &current_presentation);
         let status_changed = previous_status.is_some_and(|previous| previous != current_status);
-        let should_emit = if self.emit_initial_match {
-            self.emit_initial_match = false;
-            self.status_filter
-                .is_some_and(|wanted| wanted == current_status)
-        } else {
-            (status_changed || presentation_changed)
-                && !self
-                    .status_filter
-                    .is_some_and(|wanted| wanted != current_status)
-        };
-        if !should_emit {
+        if !(status_changed || presentation_changed) {
+            return None;
+        }
+        if self
+            .status_filter
+            .is_some_and(|wanted| wanted != current_status)
+        {
             return None;
         }
 
@@ -412,3 +524,151 @@ fn pane_get(
         },
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use super::*;
+    use crate::api::schema::{AgentStatus, EventData, EventEnvelope, EventKind};
+
+    fn status_event(custom_status: Option<&str>) -> EventEnvelope {
+        EventEnvelope {
+            event: EventKind::PaneAgentStatusChanged,
+            data: EventData::PaneAgentStatusChanged {
+                pane_id: "pane_1".into(),
+                workspace_id: "workspace_1".into(),
+                agent_status: AgentStatus::Working,
+                agent: Some("pi".into()),
+                title: None,
+                display_agent: None,
+                custom_status: custom_status.map(str::to_string),
+                state_labels: HashMap::new(),
+            },
+        }
+    }
+
+    #[test]
+    fn agent_status_subscription_replays_queued_metadata_set_and_expiry_events() {
+        let event_hub = EventHub::default();
+        let mut subscription = ActiveAgentStatusChangedSubscription {
+            pane_id: "pane_1".into(),
+            status_filter: None,
+            last_status: Some(AgentStatus::Working),
+            last_presentation: Some(PanePresentationSnapshot {
+                title: None,
+                display_agent: None,
+                custom_status: None,
+                state_labels: HashMap::new(),
+            }),
+            last_sequence: event_hub.current_sequence(),
+            initial_event: None,
+            request_prefix: "test".into(),
+        };
+
+        event_hub.push(status_event(Some("short lived")));
+        event_hub.push(status_event(None));
+
+        let set_event = subscription
+            .poll(&tokio::sync::mpsc::unbounded_channel().0, &event_hub)
+            .expect("set event");
+        let SubscriptionEventData::PaneAgentStatusChanged(set_data) = set_event.data else {
+            panic!("wrong event data");
+        };
+        assert_eq!(set_data.custom_status.as_deref(), Some("short lived"));
+
+        let expiry_event = subscription
+            .poll(&tokio::sync::mpsc::unbounded_channel().0, &event_hub)
+            .expect("expiry event");
+        let SubscriptionEventData::PaneAgentStatusChanged(expiry_data) = expiry_event.data else {
+            panic!("wrong event data");
+        };
+        assert_eq!(expiry_data.custom_status, None);
+    }
+
+    #[test]
+    fn agent_status_subscription_prefers_setup_window_events_over_initial_snapshot() {
+        let event_hub = EventHub::default();
+        let mut subscription = ActiveAgentStatusChangedSubscription {
+            pane_id: "pane_1".into(),
+            status_filter: Some(AgentStatus::Working),
+            last_status: Some(AgentStatus::Working),
+            last_presentation: Some(PanePresentationSnapshot {
+                title: None,
+                display_agent: None,
+                custom_status: None,
+                state_labels: HashMap::new(),
+            }),
+            last_sequence: event_hub.current_sequence(),
+            initial_event: Some(PaneAgentStatusChangedEvent {
+                pane_id: "pane_1".into(),
+                workspace_id: "workspace_1".into(),
+                agent_status: AgentStatus::Working,
+                agent: Some("pi".into()),
+                title: None,
+                display_agent: None,
+                custom_status: None,
+                state_labels: HashMap::new(),
+            }),
+            request_prefix: "test".into(),
+        };
+
+        event_hub.push(status_event(Some("short lived")));
+        event_hub.push(status_event(None));
+
+        let set_event = subscription
+            .poll(&tokio::sync::mpsc::unbounded_channel().0, &event_hub)
+            .expect("set event");
+        let SubscriptionEventData::PaneAgentStatusChanged(set_data) = set_event.data else {
+            panic!("wrong event data");
+        };
+        assert_eq!(set_data.custom_status.as_deref(), Some("short lived"));
+
+        let expiry_event = subscription
+            .poll(&tokio::sync::mpsc::unbounded_channel().0, &event_hub)
+            .expect("expiry event");
+        let SubscriptionEventData::PaneAgentStatusChanged(expiry_data) = expiry_event.data else {
+            panic!("wrong event data");
+        };
+        assert_eq!(expiry_data.custom_status, None);
+    }
+
+    #[test]
+    fn agent_status_subscription_emits_setup_window_event_already_reflected_by_probe() {
+        let event_hub = EventHub::default();
+        let mut subscription = ActiveAgentStatusChangedSubscription {
+            pane_id: "pane_1".into(),
+            status_filter: Some(AgentStatus::Working),
+            last_status: Some(AgentStatus::Working),
+            last_presentation: Some(PanePresentationSnapshot {
+                title: None,
+                display_agent: None,
+                custom_status: Some("short lived".into()),
+                state_labels: HashMap::new(),
+            }),
+            last_sequence: event_hub.current_sequence(),
+            initial_event: Some(PaneAgentStatusChangedEvent {
+                pane_id: "pane_1".into(),
+                workspace_id: "workspace_1".into(),
+                agent_status: AgentStatus::Working,
+                agent: Some("pi".into()),
+                title: None,
+                display_agent: None,
+                custom_status: Some("short lived".into()),
+                state_labels: HashMap::new(),
+            }),
+            request_prefix: "test".into(),
+        };
+
+        event_hub.push(status_event(Some("short lived")));
+
+        let event = subscription
+            .poll(&tokio::sync::mpsc::unbounded_channel().0, &event_hub)
+            .expect("setup-window event");
+        let SubscriptionEventData::PaneAgentStatusChanged(data) = event.data else {
+            panic!("wrong event data");
+        };
+        assert_eq!(data.custom_status.as_deref(), Some("short lived"));
+        assert!(subscription.initial_event.is_none());
+    }
+}
diff --git a/src/app/actions.rs b/src/app/actions.rs
index 9b18eef7..3dd5562f 100644
--- a/src/app/actions.rs
+++ b/src/app/actions.rs
@@ -93,10 +93,12 @@ pub struct PaneStateUpdate {
     pub previous_agent_label: Option<String>,
     pub previous_known_agent: Option<Agent>,
     pub previous_state: AgentState,
+    pub previous_seen: bool,
     pub previous_presentation: crate::terminal::EffectivePresentation,
     pub agent_label: Option<String>,
     pub known_agent: Option<Agent>,
     pub state: AgentState,
+    pub seen: bool,
     pub presentation: crate::terminal::EffectivePresentation,
 }
 
@@ -702,24 +704,27 @@ impl AppState {
         pane_terminals
             .into_iter()
             .filter_map(|(ws_idx, pane_id, terminal_id)| {
+                let previous_seen = self.workspaces[ws_idx].pane_state(pane_id)?.seen;
                 let mutation = self
                     .terminals
                     .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 update = PaneStateUpdate {
                     pane_id,
                     ws_idx,
                     previous_agent_label: change.previous_agent_label.clone(),
                     previous_known_agent: change.previous_known_agent,
                     previous_state: change.previous_state,
+                    previous_seen,
                     previous_presentation: change.previous_presentation.clone(),
                     agent_label: change.agent_label.clone(),
                     known_agent: change.known_agent,
                     state: change.state,
+                    seen,
                     presentation: change.presentation.clone(),
                 };
-                self.apply_pane_state_change(ws_idx, pane_id, &change);
                 Some(update)
             })
             .collect()
@@ -2059,6 +2064,7 @@ impl AppState {
             .pane_state(pane_id)?
             .attached_terminal_id
             .clone();
+        let previous_seen = self.workspaces[ws_idx].pane_state(pane_id)?.seen;
         let mutation = {
             let terminal = self.terminals.get_mut(&terminal_id)?;
             update(terminal)?
@@ -2067,19 +2073,21 @@ impl AppState {
             self.mark_session_dirty();
         }
         let change = mutation.effective_state_change?;
+        let seen = self.apply_pane_state_change(ws_idx, pane_id, &change)?;
         let update = PaneStateUpdate {
             pane_id,
             ws_idx,
             previous_agent_label: change.previous_agent_label.clone(),
             previous_known_agent: change.previous_known_agent,
             previous_state: change.previous_state,
+            previous_seen,
             previous_presentation: change.previous_presentation.clone(),
             agent_label: change.agent_label.clone(),
             known_agent: change.known_agent,
             state: change.state,
+            seen,
             presentation: change.presentation.clone(),
         };
-        self.apply_pane_state_change(ws_idx, pane_id, &change);
         Some(update)
     }
 
@@ -2088,23 +2096,21 @@ impl AppState {
         ws_idx: usize,
         pane_id: PaneId,
         change: &EffectiveStateChange,
-    ) {
+    ) -> Option<bool> {
         let is_active_tab = self.pane_is_in_active_tab(ws_idx, pane_id);
         let suppress_active_tab_notifications =
             active_tab_suppresses_notifications(is_active_tab, self.outer_terminal_focus);
-        let Some(pane) = self.workspaces[ws_idx]
+        let pane = self.workspaces[ws_idx]
             .tabs
             .iter_mut()
-            .find_map(|tab| tab.panes.get_mut(&pane_id))
-        else {
-            return;
-        };
+            .find_map(|tab| tab.panes.get_mut(&pane_id))?;
 
         if change.state != AgentState::Idle {
             pane.seen = true;
         } else if is_background_completion_transition(change.previous_state, change.state) {
             pane.seen = suppress_active_tab_notifications;
         }
+        let seen = pane.seen;
 
         if self.local_sound_playback && self.sound.allows(change.known_agent) {
             if let Some(sound) = notification_sound_for_state_change(
@@ -2145,6 +2151,8 @@ impl AppState {
                 });
             }
         }
+
+        Some(seen)
     }
 
     fn handle_pane_died(&mut self, pane_id: PaneId) {
diff --git a/src/app/api.rs b/src/app/api.rs
index 6319bdb6..a69dccec 100644
--- a/src/app/api.rs
+++ b/src/app/api.rs
@@ -227,16 +227,18 @@ impl App {
             });
         }
 
-        if update.previous_state != update.state
+        let previous_agent_status = pane_agent_status(update.previous_state, update.previous_seen);
+        let agent_status = self
+            .state
+            .workspaces
+            .get(update.ws_idx)
+            .and_then(|ws| ws.pane_state(update.pane_id))
+            .map(|pane| pane_agent_status(update.state, pane.seen))
+            .unwrap_or_else(|| pane_agent_status(update.state, update.seen));
+
+        if previous_agent_status != agent_status
             || update.previous_presentation != update.presentation
         {
-            let agent_status = self
-                .state
-                .workspaces
-                .get(update.ws_idx)
-                .and_then(|ws| ws.pane_state(update.pane_id))
-                .map(|pane| pane_agent_status(update.state, pane.seen))
-                .unwrap_or_else(|| pane_agent_status(update.state, true));
             let presentation = update.presentation.clone();
             self.emit_event(crate::api::schema::EventEnvelope {
                 event: crate::api::schema::EventKind::PaneAgentStatusChanged,
@@ -244,6 +246,7 @@ impl App {
                     pane_id,
                     workspace_id,
                     agent_status,
+                    agent: update.agent_label.clone(),
                     title: presentation.title,
                     display_agent: presentation.display_agent,
                     custom_status: presentation.custom_status,