mark forgotten drafts in pi agents' input boxes, and keep them across a crash or restore

5 files changed, 236 insertions, 20 deletions

diff --git a/crates/dinggerd/assets/pi/pi-agent-state.ts b/crates/dinggerd/assets/pi/pi-agent-state.ts
index aae40cc..e4e357d 100644
--- a/crates/dinggerd/assets/pi/pi-agent-state.ts
+++ b/crates/dinggerd/assets/pi/pi-agent-state.ts
@@ -111,6 +111,55 @@ function savedImportance(ctx: any): Importance {
   return "normal";
 }
 
+// The draft in the input box, as last reported: dingger marks a draft you
+// left behind. Each session's draft is also kept in a file of its own, so
+// it comes back when the session does after a crash, a restart, or
+// `pi --session`.
+let draft = "";
+const draftDir = path.join(path.dirname(socketPath), "drafts");
+// Enough to find the agent by it; the file keeps all of it.
+const reportedDraftChars = 500;
+
+function draftFile(): string | undefined {
+  return sessionFile ? path.join(draftDir, `${path.basename(sessionFile)}.draft`) : undefined;
+}
+
+function saveDraft(text: string): void {
+  const file = draftFile();
+  if (!file) {
+    return;
+  }
+  try {
+    if (text.trim()) {
+      fs.mkdirSync(draftDir, { recursive: true });
+      fs.writeFileSync(file, text);
+    } else {
+      fs.unlinkSync(file);
+    }
+  } catch {
+    // None to remove, or nowhere to keep it.
+  }
+}
+
+// The session's saved draft, into an empty input box.
+function restoreDraft(ctx: any): void {
+  const file = draftFile();
+  if (!file) {
+    return;
+  }
+  try {
+    if ((ctx.ui.getEditorText() ?? "") !== "") {
+      return;
+    }
+    const saved = fs.readFileSync(file, "utf8");
+    if (saved.trim()) {
+      ctx.ui.setEditorText(saved);
+    }
+  } catch {
+    // No saved draft.
+  }
+}
+
 function requestId(): string {
   return `${source}:${Date.now()}:${Math.random().toString(36).slice(2)}`;
 }
@@ -128,6 +177,7 @@ async function report(state: AgentState | undefined, message?: string): Promise<
       importance,
       extension: extensionVersion,
       ...(controlListening ? { control: controlPath } : {}),
+      draft: draft.slice(0, reportedDraftChars),
       seq: Date.now() * 1000,
       ...(sessionFile ? { agent_session_path: sessionFile } : {}),
     },
@@ -269,10 +319,37 @@ export default function (pi) {
     },
   });
 
-  // The latest context, for the control socket's checks.
+  // The latest context, for the control socket's checks and the input box.
   let latestCtx: any;
   let control: net.Server | undefined;
 
+  // Read the input box once typing pauses; report and save it when it changed.
+  let draftTimer: ReturnType<typeof setTimeout> | undefined;
+  let stopWatchingInput: (() => void) | undefined;
+
+  function checkDraft(): void {
+    let text: string;
+    try {
+      text = latestCtx?.ui?.getEditorText?.() ?? "";
+    } catch {
+      return;
+    }
+    if (text === draft) {
+      return;
+    }
+    draft = text;
+    saveDraft(text);
+    void report(undefined);
+  }
+
+  function checkDraftSoon(): void {
+    if (draftTimer) {
+      clearTimeout(draftTimer);
+    }
+    draftTimer = setTimeout(checkDraft, 500);
+    draftTimer.unref?.();
+  }
+
   // Run only by the control socket below, which checked that pi is idle
   // with an empty input box: extensions can reload pi only from a command.
   pi.registerCommand("dingger-reload", {
@@ -345,6 +422,17 @@ export default function (pi) {
   }
 
   pi.on("session_shutdown", () => {
+    // The input box as it is now: a quit right after typing comes before
+    // the next check.
+    if (draftTimer) {
+      checkDraft();
+    }
+    stopWatchingInput?.();
+    stopWatchingInput = undefined;
+    if (draftTimer) {
+      clearTimeout(draftTimer);
+      draftTimer = undefined;
+    }
     if (!control) {
       return;
     }
@@ -368,6 +456,16 @@ export default function (pi) {
     // A reload can replace this extension mid-run without emitting another agent_start.
     agentActive = ctx?.isIdle?.() === false;
     updateSessionFile(ctx);
+    restoreDraft(ctx);
+    // Before the first report, which carries it.
+    draft = "";
+    checkDraft();
+    stopWatchingInput?.();
+    // Every key reaches this before the input box: read the box after it.
+    stopWatchingInput = ctx.ui.onTerminalInput?.(() => {
+      checkDraftSoon();
+      return undefined;
+    });
     // Each session has its own: /new starts at normal, /resume brings its own back.
     importance = savedImportance(ctx);
     announceImportance();
@@ -379,6 +477,8 @@ export default function (pi) {
     agentActive = true;
     updateSessionFile(ctx);
     publishState();
+    // The prompt left the input box.
+    checkDraftSoon();
   });
 
   pi.on("agent_settled", (_event, ctx) => {
@@ -389,5 +489,7 @@ export default function (pi) {
 
     agentActive = false;
     publishState();
+    // 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 0074544..072735f 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.
 //!
+//! It also carries the `draft` in the agent's input box ("" for none), for
+//! marking a draft you left behind; never logged.
+//!
 //! Each line is answered with `{"ok":true}`.
 //!
 //! `{"method":"poke"}` only wakes the watcher; `dingger focus` and tmux's
@@ -69,6 +72,8 @@ struct PaneReports {
     /// The version of the extension that reported, and its control socket.
     extension: Option<String>,
     control: Option<String>,
+    /// The draft in the agent's input box, as one line ("" for none).
+    draft: Option<String>,
     /// When the last report arrived.
     at: Option<Instant>,
 }
@@ -165,6 +170,8 @@ struct HookParams {
     extension: Option<String>,
     #[serde(default)]
     control: Option<String>,
+    #[serde(default)]
+    draft: Option<String>,
 }
 
 fn handle_stream(stream: UnixStream) {
@@ -231,6 +238,7 @@ fn handle_stream(stream: UnixStream) {
             importance,
             extension: params.extension,
             control: params.control,
+            draft: params.draft,
         };
         record(params.pane_id, state, reported);
         wake();
@@ -244,6 +252,22 @@ struct Reported {
     importance: Option<Importance>,
     extension: Option<String>,
     control: Option<String>,
+    draft: Option<String>,
+}
+
+/// The longest draft kept: enough to find an agent by it.
+const DRAFT_CHARS: usize = 200;
+
+/// A draft as one line, at most [`DRAFT_CHARS`] long; "" for none.
+fn draft_line(draft: &str) -> String {
+    draft
+        .split(|c: char| c.is_whitespace() || c.is_control())
+        .filter(|word| !word.is_empty())
+        .collect::<Vec<_>>()
+        .join(" ")
+        .chars()
+        .take(DRAFT_CHARS)
+        .collect()
 }
 
 /// A value fit for a tmux option: not empty, no control characters.
@@ -271,6 +295,9 @@ fn record(pane_id: String, state: Option<AgentState>, reported: Reported) {
     if let Some(control) = option_value(reported.control) {
         pane.control = Some(control);
     }
+    if let Some(draft) = reported.draft {
+        pane.draft = Some(draft_line(&draft));
+    }
 }
 
 /// Start listening; reports accumulate for the watcher to consume.
@@ -328,6 +355,12 @@ pub fn extension(pane_id: &str) -> (Option<String>, Option<String>) {
     reports().get(pane_id).map(|pane| (pane.extension.clone(), pane.control.clone())).unwrap_or_default()
 }
 
+/// The draft in the pane's input box ("" for none), if its extension
+/// reported one.
+pub fn draft(pane_id: &str) -> Option<String> {
+    reports().get(pane_id).and_then(|pane| pane.draft.clone())
+}
+
 /// Ask the extension behind `control` to reload; its answer: `Ok(())` when
 /// it does, else why not (busy, a draft, or no answer).
 pub fn request_reload(control: &str) -> Result<(), String> {
@@ -408,6 +441,21 @@ mod tests {
     }
 
     #[test]
+    fn a_draft_is_kept_as_one_short_line_until_reported_again() {
+        let draft = |text: &str| Reported { draft: Some(text.into()), ..Default::default() };
+        record("%t5".into(), Some(AgentState::Idle), Reported::default());
+        assert_eq!(super::draft("%t5"), None, "an extension that does not report drafts");
+        record("%t5".into(), None, draft("  fix the\n\tsidecar;\x1b[1m  "));
+        assert_eq!(super::draft("%t5").as_deref(), Some("fix the sidecar; [1m"));
+        record("%t5".into(), Some(AgentState::Working), Reported::default());
+        assert_eq!(super::draft("%t5").as_deref(), Some("fix the sidecar; [1m"), "kept by a report without one");
+        record("%t5".into(), None, draft(&"x".repeat(500)));
+        assert_eq!(super::draft("%t5").map(|d| d.chars().count()), Some(DRAFT_CHARS));
+        record("%t5".into(), None, draft(""));
+        assert_eq!(super::draft("%t5").as_deref(), Some(""), "cleared");
+    }
+
+    #[test]
     fn a_reload_answer_says_whether_it_reloads() {
         let dir = std::env::temp_dir().join(format!("dingger-reload-{}", std::process::id()));
         std::fs::create_dir_all(&dir).unwrap();
diff --git a/crates/dinggerd/src/main.rs b/crates/dinggerd/src/main.rs
index 1fc5cd2..44db683 100644
--- a/crates/dinggerd/src/main.rs
+++ b/crates/dinggerd/src/main.rs
@@ -443,7 +443,8 @@ fn mark_unread(pane: Option<String>) -> Result<(), String> {
 }
 
 /// Every agent pane in herdr's priority order (blocked, done, working,
-/// idle), oldest-waiting first within a status.
+/// idle), oldest-waiting first within a status. ✎ marks a draft left in
+/// the agent's input box.
 fn status() -> Result<(), String> {
     let mut panes: Vec<_> = tmux::list_panes()?
         .into_iter()
@@ -456,7 +457,8 @@ fn status() -> Result<(), String> {
     panes.sort_by_key(|pane| {
         let status = watcher::Status::from_label(&pane.written_state);
         let importance = importance::Importance::from_option(&pane.written_importance);
-        (watcher::list_order(status, importance, pane.written_since.parse().ok()), pane.display_target.clone())
+        let since = pane.written_since.parse().ok();
+        (watcher::list_order(status, watcher::forgotten_draft(pane), importance, since), pane.display_target.clone())
     });
     let place = |pane: &tmux::Pane| pane.display_target.replacen(':', " · ", 1);
     let width = panes.iter().map(|pane| place(pane).chars().count()).max().unwrap_or(5).max(5);
@@ -464,12 +466,13 @@ fn status() -> Result<(), String> {
     for pane in &panes {
         let state = if pane.written_state.is_empty() { "-" } else { &pane.written_state };
         let agent = pane.agent.map(detect_core::agent_label).unwrap_or("-");
+        let draft = if watcher::forgotten_draft(pane) { "✎ " } else { "" };
         let name = match importance::Importance::from_option(&pane.written_importance) {
             importance::Importance::High => format!("! {}", pane.written_name),
             importance::Importance::Low => format!("{} (low)", pane.written_name),
             importance::Importance::Normal => pane.written_name.clone(),
         };
-        println!("{state:<8} {:<width$}  {agent:<7} {}", place(pane), name.trim());
+        println!("{state:<8} {:<width$}  {agent:<7} {draft}{}", place(pane), name.trim());
     }
     Ok(())
 }
diff --git a/crates/dinggerd/src/tmux.rs b/crates/dinggerd/src/tmux.rs
index b743256..510c54b 100644
--- a/crates/dinggerd/src/tmux.rs
+++ b/crates/dinggerd/src/tmux.rs
@@ -32,6 +32,8 @@ pub struct Pane {
     pub written_importance: String,
     pub written_extension: String,
     pub written_control: String,
+    /// The draft in the agent's input box, as one line.
+    pub written_draft: String,
 }
 
 
@@ -115,14 +117,15 @@ pub fn list_panes() -> Result<Vec<Pane>, String> {
         "#{@agent_session}\t",
         "#{@dingger_importance}\t",
         "#{@dingger_extension}\t",
-        "#{@dingger_control}",
+        "#{@dingger_control}\t",
+        "#{@dingger_draft}",
     );
     let out = run_tmux(&["list-panes", "-a", "-F", format])?;
     let mut panes: Vec<Pane> = Vec::new();
     let mut process_tree = None;
     for line in out.lines() {
         let fields: Vec<&str> = line.split('\t').collect();
-        if fields.len() < 22 {
+        if fields.len() < 23 {
             continue;
         }
         let attended = fields[4] == "1" && fields[5] != "0";
@@ -167,6 +170,7 @@ pub fn list_panes() -> Result<Vec<Pane>, String> {
             written_importance: fields[19].to_string(),
             written_extension: fields[20].to_string(),
             written_control: fields[21].to_string(),
+            written_draft: fields[22].to_string(),
         });
     }
     Ok(panes)
@@ -242,20 +246,33 @@ pub fn set_pane_option(pane_id: &str, name: &str, value: &str) -> Result<(), Str
 /// Set (`Some`) or unset (`None`) several pane options in one tmux command
 /// list, applied as a unit.
 pub fn set_pane_options(pane_id: &str, changes: &[(&str, Option<String>)]) -> Result<(), String> {
-    let mut args: Vec<&str> = Vec::new();
+    let mut args: Vec<String> = Vec::new();
     for (name, value) in changes {
         if !args.is_empty() {
-            args.push(";");
+            args.push(";".into());
         }
         match value {
-            Some(value) => args.extend(["set-option", "-p", "-t", pane_id, name, value.as_str()]),
-            None => args.extend(["set-option", "-p", "-u", "-t", pane_id, name]),
+            Some(value) => {
+                args.extend(["set-option", "-p", "-t", pane_id, name].map(String::from));
+                args.push(command_list_value(value));
+            }
+            None => args.extend(["set-option", "-p", "-u", "-t", pane_id, name].map(String::from)),
         }
     }
     if args.is_empty() {
         return Ok(());
     }
-    run_tmux(&args).map(|_| ())
+    run_tmux(&args.iter().map(String::as_str).collect::<Vec<_>>()).map(|_| ())
+}
+
+/// A value as an argument in a tmux command list: tmux ends a command at an
+/// argument ending in `;` (dropping it) unless it ends in `\;`, and free
+/// text (a name, a draft) may end in one.
+fn command_list_value(value: &str) -> String {
+    match value.strip_suffix(';') {
+        Some(rest) => format!("{rest}\\;"),
+        None => value.to_string(),
+    }
 }
 
 pub fn unset_pane_option(pane_id: &str, name: &str) -> Result<(), String> {
diff --git a/crates/dinggerd/src/watcher.rs b/crates/dinggerd/src/watcher.rs
index 1c57609..b0acc1d 100644
--- a/crates/dinggerd/src/watcher.rs
+++ b/crates/dinggerd/src/watcher.rs
@@ -26,6 +26,7 @@
 //! - `@agent_session` the agent's session reference, for restore
 //! - `@dingger_extension`, `@dingger_control` the version of pi's extension
 //!   and its control socket, for reloading agents on an old one
+//! - `@dingger_draft` the draft in the agent's input box, as one line
 
 use std::collections::HashMap;
 
@@ -194,21 +195,36 @@ pub fn mark_unread(track: PaneTrack, attended: bool, now: u64) -> PaneTrack {
     PaneTrack { seen: false, held: attended, since: Some(now), ..track }
 }
 
-/// The order agents are listed in: herdr's status priority; within a
-/// status, the more important first; then done and blocked oldest first
-/// (the queue: nothing waits forever), working and idle most recent first;
-/// then by place.
+/// The order agents are listed in: herdr's status priority; idle agents
+/// with a forgotten draft before the other idle ones; within a status, the
+/// more important first; then done and blocked oldest first (the queue:
+/// nothing waits forever), working and idle most recent first; then by
+/// place.
 pub fn list_order(
     status: Status,
+    forgotten_draft: bool,
     importance: Importance,
     since: Option<u64>,
-) -> (std::cmp::Reverse<u8>, std::cmp::Reverse<Importance>, i128) {
+) -> (std::cmp::Reverse<u8>, bool, std::cmp::Reverse<Importance>, i128) {
     let since = since.map(i128::from);
     let key = match status {
         Status::Done | Status::Blocked => since.unwrap_or(i128::MAX),
         _ => since.map(|since| -since).unwrap_or(i128::MAX),
     };
-    (std::cmp::Reverse(status.priority()), std::cmp::Reverse(importance), key)
+    let draft_last = !(forgotten_draft && status == Status::Idle);
+    (std::cmp::Reverse(status.priority()), draft_last, std::cmp::Reverse(importance), key)
+}
+
+/// A draft left in an agent's input box while the user looks elsewhere
+/// (where they look, they may be typing it).
+pub fn forgotten_draft(pane: &tmux::Pane) -> bool {
+    pane.agent.is_some() && !pane.attended && !draft_of(pane).is_empty()
+}
+
+/// The pane's draft: as last reported, else as written (it survives a
+/// daemon restart).
+fn draft_of(pane: &tmux::Pane) -> String {
+    hooks::draft(&pane.pane_id).unwrap_or_else(|| pane.written_draft.clone())
 }
 
 /// A pane's importance: as its agent last reported it, else as written
@@ -493,6 +509,12 @@ impl Watcher {
                 change("@dingger_control", &pane.written_control, String::new());
             }
         }
+        // The draft goes with the agent too.
+        match (pane.agent, hooks::draft(&pane.pane_id)) {
+            (Some(_), Some(draft)) => change("@dingger_draft", &pane.written_draft, draft),
+            (None, _) => change("@dingger_draft", &pane.written_draft, String::new()),
+            _ => {}
+        }
         if let Err(err) = tmux::set_pane_options(&pane.pane_id, &changes) {
             warn!(target = %pane.display_target, %err, "failed to write pane options");
         }
@@ -539,7 +561,10 @@ impl Watcher {
             }
             // The queue: what needs the user, in list order.
             if matches!(track.status(), Status::Blocked | Status::Done) {
-                queue.push((list_order(track.status(), importance_of(pane), track.since), pane.pane_id.as_str()));
+                queue.push((
+                    list_order(track.status(), forgotten_draft(pane), importance_of(pane), track.since),
+                    pane.pane_id.as_str(),
+                ));
             }
         }
         queue.sort();
@@ -801,7 +826,7 @@ mod tests {
             ("blocked", Status::Blocked, Some(99)),
             ("working new", Status::Working, Some(70)),
         ];
-        rows.sort_by_key(|(_, status, since)| list_order(*status, Importance::Normal, *since));
+        rows.sort_by_key(|(_, status, since)| list_order(*status, false, Importance::Normal, *since));
         let names: Vec<_> = rows.iter().map(|(name, ..)| *name).collect();
         assert_eq!(
             names,
@@ -820,7 +845,7 @@ mod tests {
             ("working high", Status::Working, High, Some(1)),
             ("working new", Status::Working, Normal, Some(80)),
         ];
-        rows.sort_by_key(|(_, status, importance, since)| list_order(*status, *importance, *since));
+        rows.sort_by_key(|(_, status, importance, since)| list_order(*status, false, *importance, *since));
         let names: Vec<_> = rows.iter().map(|(name, ..)| *name).collect();
         assert_eq!(
             names,
@@ -830,6 +855,27 @@ mod tests {
     }
 
     #[test]
+    fn a_forgotten_draft_lifts_an_idle_agent_only() {
+        use Importance::*;
+        let mut rows = [
+            ("idle new", Status::Idle, false, High, Some(90)),
+            ("idle draft old low", Status::Idle, true, Low, Some(5)),
+            ("idle draft new", Status::Idle, true, Normal, Some(50)),
+            ("done old", Status::Done, false, Normal, Some(10)),
+            ("done draft new", Status::Done, true, Normal, Some(60)),
+            ("working draft", Status::Working, true, Normal, Some(1)),
+            ("working new", Status::Working, false, Normal, Some(80)),
+        ];
+        rows.sort_by_key(|(_, status, draft, importance, since)| list_order(*status, *draft, *importance, *since));
+        let names: Vec<_> = rows.iter().map(|(name, ..)| *name).collect();
+        assert_eq!(
+            names,
+            ["done old", "done draft new", "working new", "working draft", "idle draft new", "idle draft old low", "idle new"],
+            "status first; among idle agents, a draft before importance"
+        );
+    }
+
+    #[test]
     fn agent_exit_is_unknown() {
         let done = step(step(PaneTrack::NEW, Working, false), Idle, false);
         assert_eq!(step(done, Unknown, false).status(), Status::Unknown);