pi: a dialog during a run is blocked, and a cancelled run is idle, not done

3 files changed, 95 insertions, 21 deletions

diff --git a/crates/dinggerd/assets/pi/pi-agent-state.ts b/crates/dinggerd/assets/pi/pi-agent-state.ts
index fc70638..62f29db 100644
--- a/crates/dinggerd/assets/pi/pi-agent-state.ts
+++ b/crates/dinggerd/assets/pi/pi-agent-state.ts
@@ -167,7 +167,7 @@ function requestId(): string {
   return `${source}:${Date.now()}:${Math.random().toString(36).slice(2)}`;
 }
 
-async function report(state: AgentState | undefined, message?: string): Promise<void> {
+async function report(state: AgentState | undefined, message?: string, aborted = false): Promise<void> {
   await sendRequest({
     id: requestId(),
     method: "pane.report_agent",
@@ -177,6 +177,8 @@ async function report(state: AgentState | undefined, message?: string): Promise<
       agent: "pi",
       ...(state ? { state } : {}),
       message,
+      // The run this idle report ends was cancelled: not done.
+      ...(aborted ? { aborted: true } : {}),
       importance,
       extension: extensionVersion,
       ...(controlListening ? { control: controlPath } : {}),
@@ -189,10 +191,10 @@ async function report(state: AgentState | undefined, message?: string): Promise<
 }
 
 let sendInFlight = false;
-let queuedState: { state: AgentState; message?: string } | undefined;
+let queuedState: { state: AgentState; message?: string; aborted: boolean } | undefined;
 
-function queueState(state: AgentState, message?: string): void {
-  queuedState = { state, message };
+function queueState(state: AgentState, message?: string, aborted = false): void {
+  queuedState = { state, message, aborted };
   if (!sendInFlight) {
     void drainStateQueue();
   }
@@ -208,7 +210,7 @@ async function drainStateQueue(): Promise<void> {
     while (queuedState) {
       const next = queuedState;
       queuedState = undefined;
-      await report(next.state, next.message);
+      await report(next.state, next.message, next.aborted);
     }
   } finally {
     sendInFlight = false;
@@ -239,16 +241,42 @@ export default function (pi) {
     return { state: "idle" as const, message: undefined };
   }
 
-  function publishState(force = false) {
+  function publishState(force = false, aborted = false) {
     const next = desiredState();
     if (!force && next.state === lastState && next.message === lastMessage) {
       return;
     }
     lastState = next.state;
     lastMessage = next.message;
-    queueState(next.state, next.message);
+    queueState(next.state, next.message, aborted && next.state === "idle");
   }
 
+  // A dialog that opens during a run (an extension asking before it goes
+  // on): the agent waits on you, so it is blocked until the dialog closes.
+  // One opened while idle comes from what you just did, and changes nothing.
+  let dialogsBlocking = 0;
+  pi.on("ui_prompt_start", (event, ctx) => {
+    latestCtx = ctx;
+    if (!agentActive) {
+      return;
+    }
+    dialogsBlocking += 1;
+    blockedCount += 1;
+    blockedMessage = event?.title;
+    publishState();
+  });
+  pi.on("ui_prompt_end", () => {
+    if (dialogsBlocking === 0) {
+      return;
+    }
+    dialogsBlocking -= 1;
+    blockedCount = Math.max(0, blockedCount - 1);
+    if (blockedCount === 0) {
+      blockedMessage = undefined;
+    }
+    publishState();
+  });
+
   pi.events.on("blocked", (data) => {
     if (!data?.active) {
       blockedCount = Math.max(0, blockedCount - 1);
@@ -516,7 +544,7 @@ export default function (pi) {
     });
   }
 
-  pi.on("agent_settled", (_event, ctx) => {
+  pi.on("agent_settled", (event, ctx) => {
     latestCtx = ctx;
     if (ctx?.isIdle?.() !== true) {
       return;
@@ -525,7 +553,8 @@ export default function (pi) {
     agentActive = false;
     // The last reply's usage: message_end comes before pi counts it.
     checkContext(ctx);
-    publishState();
+    // A cancelled run (Escape) is idle, not done.
+    publishState(false, event?.aborted === true);
     // Escape puts queued messages back into the input box.
     checkDraftSoon();
   });
diff --git a/crates/dinggerd/src/hooks.rs b/crates/dinggerd/src/hooks.rs
index 82fa7bd..709f62b 100644
--- a/crates/dinggerd/src/hooks.rs
+++ b/crates/dinggerd/src/hooks.rs
@@ -27,6 +27,9 @@
 //! `{"method":"agents.reload_outdated"}`, and the watcher reloads each agent
 //! on an old extension once it is idle and off screen.
 //!
+//! An idle report may say the run was `aborted` (cancelled, as with
+//! Escape): the agent is idle, not done, and nobody is told it finished.
+//!
 //! It also carries the `draft` in the agent's input box ("" for none), for
 //! marking a draft you left behind; never logged. And `context`: how full
 //! the agent's context window is, in percent, as pi's footer shows it (""
@@ -66,7 +69,7 @@ struct PaneReports {
     /// Last reported state: authoritative until the next report.
     last: Option<AgentState>,
     /// Reports not yet applied by the watcher, oldest first.
-    pending: Vec<AgentState>,
+    pending: Vec<Observation>,
     /// The agent's latest session reference.
     session: Option<String>,
     /// The agent's importance, if reported.
@@ -178,6 +181,8 @@ struct HookParams {
     draft: Option<String>,
     #[serde(default)]
     context: Option<String>,
+    #[serde(default)]
+    aborted: Option<bool>,
 }
 
 fn handle_stream(stream: UnixStream) {
@@ -246,6 +251,7 @@ fn handle_stream(stream: UnixStream) {
             control: params.control,
             draft: params.draft,
             context: params.context,
+            aborted: params.aborted == Some(true),
         };
         record(params.pane_id, state, reported);
         wake();
@@ -261,6 +267,22 @@ struct Reported {
     control: Option<String>,
     draft: Option<String>,
     context: Option<String>,
+    /// The run this report ends was cancelled.
+    aborted: bool,
+}
+
+/// A reported state, and whether it ends a run the user cancelled.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct Observation {
+    pub state: AgentState,
+    /// Idle after a cancelled run: not done, and no "finished" alert.
+    pub cancelled: bool,
+}
+
+impl Observation {
+    pub fn plain(state: AgentState) -> Self {
+        Observation { state, cancelled: false }
+    }
 }
 
 /// The longest draft kept: enough to find an agent by it.
@@ -289,7 +311,7 @@ fn record(pane_id: String, state: Option<AgentState>, reported: Reported) {
     pane.at = Some(Instant::now());
     if let Some(state) = state {
         pane.last = Some(state);
-        pane.pending.push(state);
+        pane.pending.push(Observation { state, cancelled: reported.aborted && state == AgentState::Idle });
     }
     if let Some(session) = option_value(reported.session) {
         pane.session = Some(session);
@@ -352,11 +374,11 @@ pub fn spawn_listener() -> Result<(), String> {
 /// The pane's hook observations for this tick: every report since the last
 /// tick, in order; otherwise the last reported state; `None` if the pane
 /// never reported to this process.
-pub fn take_observations(pane_id: &str) -> Option<Vec<AgentState>> {
+pub fn take_observations(pane_id: &str) -> Option<Vec<Observation>> {
     let mut reports = reports();
     let pane = reports.get_mut(pane_id)?;
     if pane.pending.is_empty() {
-        return pane.last.map(|state| vec![state]);
+        return pane.last.map(|state| vec![Observation::plain(state)]);
     }
     Some(std::mem::take(&mut pane.pending))
 }
@@ -467,6 +489,24 @@ mod tests {
     }
 
     #[test]
+    fn a_cancelled_run_ends_in_a_cancelled_idle_observation() {
+        let aborted = || Reported { aborted: true, ..Default::default() };
+        record("%t7".into(), Some(AgentState::Working), aborted());
+        record("%t7".into(), Some(AgentState::Idle), aborted());
+        assert_eq!(
+            take_observations("%t7"),
+            Some(vec![
+                Observation::plain(AgentState::Working),
+                Observation { state: AgentState::Idle, cancelled: true },
+            ]),
+            "only an idle report ends a run"
+        );
+        assert_eq!(take_observations("%t7"), Some(vec![Observation::plain(AgentState::Idle)]), "seen once");
+        record("%t7".into(), Some(AgentState::Idle), Reported::default());
+        assert_eq!(take_observations("%t7"), Some(vec![Observation::plain(AgentState::Idle)]));
+    }
+
+    #[test]
     fn context_use_is_kept_as_whole_percent_until_reported_again() {
         let context = |text: &str| Reported { context: Some(text.into()), ..Default::default() };
         assert_eq!(super::context("%t6"), None, "an extension that does not report it");
diff --git a/crates/dinggerd/src/watcher.rs b/crates/dinggerd/src/watcher.rs
index 77b9259..aee8d5e 100644
--- a/crates/dinggerd/src/watcher.rs
+++ b/crates/dinggerd/src/watcher.rs
@@ -35,7 +35,7 @@ use std::collections::HashMap;
 use detect_core::{detect_agent_with_osc, Agent, AgentState};
 use tracing::{debug, info, warn};
 
-use crate::hooks;
+use crate::hooks::{self, Observation};
 use crate::importance::Importance;
 use crate::notify::{play_sound, send_notification, NotificationKind};
 use crate::tmux;
@@ -376,7 +376,12 @@ impl Watcher {
             // finished between two ticks is still a completion.
             let mut alerts = Vec::new();
             for observed in observations {
-                let next = advance(track, observed, pane.attended, acked, now);
+                let mut next = advance(track, observed.state, pane.attended, acked, now);
+                // A cancelled run did not finish: idle, as if seen.
+                if observed.cancelled && next.state == AgentState::Idle {
+                    next.seen = true;
+                    next.held = false;
+                }
                 if next.state != track.state {
                     info!(
                         target = %pane.display_target,
@@ -386,7 +391,7 @@ impl Watcher {
                     );
                     // Notify on transitions observed by this process, not
                     // when resuming from options written before a restart.
-                    if !restored {
+                    if !restored && !observed.cancelled {
                         alerts.extend(notification_kind_for_transition(track.state, next.state));
                     }
                 }
@@ -615,14 +620,14 @@ impl Watcher {
 
 /// The pane's raw agent states this tick, in order, or `None` when the
 /// screen is a viewer (transcript, picker) that says nothing about state.
-fn observe(pane: &tmux::Pane, listed: std::time::Instant) -> Option<Vec<AgentState>> {
+fn observe(pane: &tmux::Pane, listed: std::time::Instant) -> Option<Vec<Observation>> {
     if pane.dead {
-        return Some(vec![AgentState::Idle]);
+        return Some(vec![Observation::plain(AgentState::Idle)]);
     }
     let Some(agent) = pane.agent else {
         // The agent is gone; its last hook report must not outlive it.
         hooks::drop_reports(&pane.pane_id, listed);
-        return Some(vec![AgentState::Unknown]);
+        return Some(vec![Observation::plain(AgentState::Unknown)]);
     };
     // Agent-reported state is authoritative; it survives restarts through
     // the `@dingger_hook` option until the extension reports again.
@@ -630,7 +635,7 @@ fn observe(pane: &tmux::Pane, listed: std::time::Instant) -> Option<Vec<AgentSta
         return Some(states);
     }
     if let Some(state) = parse_state(&pane.written_hook) {
-        return Some(vec![state]);
+        return Some(vec![Observation::plain(state)]);
     }
     let screen = match tmux::capture_pane(&pane.pane_id) {
         Ok(screen) => screen,
@@ -642,7 +647,7 @@ fn observe(pane: &tmux::Pane, listed: std::time::Instant) -> Option<Vec<AgentSta
     // tmux pane_title carries the last OSC title the application set; tmux
     // does not re-expose OSC progress sequences.
     let detection = detect_agent_with_osc(Some(agent), &screen, &pane.pane_title, "");
-    (!detection.skip_state_update).then(|| vec![detection.state])
+    (!detection.skip_state_update).then(|| vec![Observation::plain(detection.state)])
 }
 
 /// The host name tmux uses as a pane's default title.