feat: let agents report their own resume command (#4687)

* feat: let agents report their own resume command

* fix: keep reported resume commands tied to accepted sessions and current owners

2 files changed, 210 insertions, 8 deletions

diff --git a/src/persist/restore.rs b/src/persist/restore.rs
index 265b4874..e79a0f3b 100644
--- a/src/persist/restore.rs
+++ b/src/persist/restore.rs
@@ -15,7 +15,8 @@ use crate::terminal::{TerminalId, TerminalRuntime, TerminalState};
 use crate::workspace::Workspace;
 
 use super::snapshot::{
-    PaneAgentSessionSnapshot, PaneHistorySnapshot, TabHistorySnapshot, WorkspaceHistorySnapshot,
+    PaneAgentResumeSnapshot, PaneAgentSessionSnapshot, PaneHistorySnapshot, TabHistorySnapshot,
+    WorkspaceHistorySnapshot,
 };
 use super::{
     DirectionSnapshot, LayoutSnapshot, SessionHistorySnapshot, SessionSnapshot, TabSnapshot,
@@ -455,6 +456,9 @@ fn unavailable_restored_terminal(
         if let Some(session) = restored_terminal_agent_session(pane.agent_session.as_ref(), false) {
             terminal.set_persisted_agent_session(session);
         }
+        if let Some(resume) = saved_reported_resume(pane) {
+            terminal.restore_reported_resume(reported_resume_from_snapshot(resume));
+        }
         match (
             pane.agent_name.as_ref(),
             pane.managed_agent_kind
@@ -530,6 +534,7 @@ fn restore_tab(
             .and_then(crate::detect::parse_canonical_agent_label);
         let saved_launch_argv = saved_pane.and_then(|p| p.launch_argv.clone());
         let saved_agent_session = saved_pane.and_then(|p| p.agent_session.as_ref());
+        let saved_agent_resume = saved_pane.and_then(saved_reported_resume);
         let saved_history =
             old_id.and_then(|old_id| history.and_then(|history| history.panes.get(old_id)));
         let startup = {
@@ -537,8 +542,17 @@ fn restore_tab(
                 enabled: runtime_context.resume_agents_on_restore,
                 resumed_sessions: resumed_agent_sessions,
             };
-            pane_restore_startup(saved_agent_session, saved_history, &mut agent_restore)
+            pane_restore_startup(
+                saved_agent_session,
+                saved_agent_resume,
+                &cwd,
+                saved_history,
+                &mut agent_restore,
+            )
         };
+        let restored_agent_resume = saved_agent_resume
+            .filter(|_| !startup.duplicate_agent_session)
+            .map(reported_resume_from_snapshot);
         let restored_agent_session =
             restored_terminal_agent_session(saved_agent_session, startup.duplicate_agent_session);
         let initial_restore_agent = startup
@@ -580,6 +594,9 @@ fn restore_tab(
             if let Some(session) = restored_agent_session {
                 terminal.set_persisted_agent_session(session);
             }
+            if let Some(resume) = restored_agent_resume {
+                terminal.restore_reported_resume(resume);
+            }
             match (saved_agent_name, saved_managed_agent) {
                 (Some(agent_name), Some(agent)) => {
                     terminal.restore_managed_agent(agent_name, agent)
@@ -677,6 +694,9 @@ fn restore_tab(
                 if let Some(session) = restored_agent_session {
                     terminal.set_persisted_agent_session(session);
                 }
+                if let Some(resume) = restored_agent_resume {
+                    terminal.restore_reported_resume(resume);
+                }
                 match (saved_agent_name, saved_managed_agent) {
                     (Some(agent_name), Some(agent)) if was_imported => {
                         terminal.restore_managed_agent(agent_name, agent)
@@ -787,6 +807,8 @@ fn restore_tab(
 
 fn pane_restore_startup<'a>(
     session: Option<&PaneAgentSessionSnapshot>,
+    reported_resume: Option<&PaneAgentResumeSnapshot>,
+    cwd: &std::path::Path,
     history: Option<&'a PaneHistorySnapshot>,
     agent_restore: &mut AgentRestoreState<'_>,
 ) -> PaneRestoreStartup<'a> {
@@ -794,8 +816,13 @@ fn pane_restore_startup<'a>(
     // resumable agent session and resume is enabled, do not replay saved pane
     // presentation history into that terminal, even when this pane is a
     // duplicate suppressed by session de-duplication.
-    let restore_plan =
-        session.and_then(|session| restore_plan_for_snapshot(session, agent_restore.enabled));
+    let restore_plan = if !agent_restore.enabled {
+        None
+    } else if let Some(resume) = reported_resume {
+        Some(reported_resume_from_snapshot(resume).plan(cwd))
+    } else {
+        session.and_then(|session| restore_plan_for_snapshot(session, true))
+    };
     let has_native_agent_restore = restore_plan.is_some();
     // Reserve before spawning so later panes in the same restore pass cannot
     // launch the same native agent session. The caller rolls this reservation
@@ -830,6 +857,22 @@ fn pane_restore_startup<'a>(
     }
 }
 
+fn saved_reported_resume(pane: &super::snapshot::PaneSnapshot) -> Option<&PaneAgentResumeSnapshot> {
+    pane.agent_resume
+        .as_ref()
+        .filter(|resume| crate::agent_resume::validate_resume_argv(&resume.argv).is_ok())
+}
+
+fn reported_resume_from_snapshot(
+    resume: &PaneAgentResumeSnapshot,
+) -> crate::agent_resume::ReportedAgentResume {
+    crate::agent_resume::ReportedAgentResume {
+        source: resume.source.clone(),
+        agent: resume.agent.clone(),
+        argv: resume.argv.clone(),
+    }
+}
+
 fn restore_plan_for_snapshot(
     session: &PaneAgentSessionSnapshot,
     resume_agents_on_restore: bool,
@@ -1124,7 +1167,13 @@ mod tests {
             resumed_sessions: &mut resumed,
         };
 
-        let startup = pane_restore_startup(Some(&session), Some(&history), &mut agent_restore);
+        let startup = pane_restore_startup(
+            Some(&session),
+            None,
+            std::path::Path::new("/a"),
+            Some(&history),
+            &mut agent_restore,
+        );
 
         assert!(startup.restore_plan.is_some());
         assert!(startup.initial_history_ansi.is_none());
@@ -1149,8 +1198,20 @@ mod tests {
             resumed_sessions: &mut resumed,
         };
 
-        let first = pane_restore_startup(Some(&session), Some(&history), &mut agent_restore);
-        let duplicate = pane_restore_startup(Some(&session), Some(&history), &mut agent_restore);
+        let first = pane_restore_startup(
+            Some(&session),
+            None,
+            std::path::Path::new("/a"),
+            Some(&history),
+            &mut agent_restore,
+        );
+        let duplicate = pane_restore_startup(
+            Some(&session),
+            None,
+            std::path::Path::new("/a"),
+            Some(&history),
+            &mut agent_restore,
+        );
 
         assert!(first.restore_plan.is_some());
         assert!(first.initial_history_ansi.is_none());
@@ -1160,6 +1221,79 @@ mod tests {
     }
 
     #[test]
+    fn pane_restore_startup_prefers_reported_resume_argv() {
+        let session = super::super::snapshot::PaneAgentSessionSnapshot {
+            source: "herdr:pi".into(),
+            agent: "pi".into(),
+            kind: crate::agent_resume::AgentSessionRefKind::Path,
+            value: test_session_path("pi-session.jsonl"),
+        };
+        let resume = super::super::snapshot::PaneAgentResumeSnapshot {
+            source: "herdr:pi".into(),
+            agent: "pi".into(),
+            argv: vec![
+                "pi".into(),
+                "--continue".into(),
+                "--model".into(),
+                "m".into(),
+            ],
+        };
+        let history = super::super::snapshot::PaneHistorySnapshot {
+            ansi: "RESTORED_HISTORY\r\n".into(),
+            lines: 1,
+        };
+        let mut resumed = HashSet::new();
+        let mut agent_restore = AgentRestoreState {
+            enabled: true,
+            resumed_sessions: &mut resumed,
+        };
+
+        let project_a = std::path::Path::new("/project-a");
+        let project_b = std::path::Path::new("/project-b");
+        let startup = pane_restore_startup(
+            Some(&session),
+            Some(&resume),
+            project_a,
+            Some(&history),
+            &mut agent_restore,
+        );
+        let plan = startup.restore_plan.expect("reported resume plan");
+        assert_eq!(plan.agent, "pi");
+        assert_eq!(plan.argv, resume.argv);
+        assert!(startup.initial_history_ansi.is_none());
+
+        let custom = super::super::snapshot::PaneAgentResumeSnapshot {
+            source: "prime-agent".into(),
+            agent: "prime-agent".into(),
+            argv: vec!["prime-agent".into(), "--continue".into()],
+        };
+        let first = pane_restore_startup(None, Some(&custom), project_a, None, &mut agent_restore);
+        assert_eq!(first.restore_plan.unwrap().argv, custom.argv);
+        let other_project =
+            pane_restore_startup(None, Some(&custom), project_b, None, &mut agent_restore);
+        assert!(other_project.restore_plan.is_some());
+        let duplicate =
+            pane_restore_startup(None, Some(&custom), project_a, None, &mut agent_restore);
+        assert!(duplicate.restore_plan.is_none());
+        assert!(duplicate.duplicate_agent_session);
+
+        let mut resumed = HashSet::new();
+        let mut disabled = AgentRestoreState {
+            enabled: false,
+            resumed_sessions: &mut resumed,
+        };
+        let startup = pane_restore_startup(
+            None,
+            Some(&custom),
+            project_a,
+            Some(&history),
+            &mut disabled,
+        );
+        assert!(startup.restore_plan.is_none());
+        assert_eq!(startup.initial_history_ansi, Some("RESTORED_HISTORY\r\n"));
+    }
+
+    #[test]
     fn pane_restore_startup_keeps_history_without_native_agent_resume() {
         let session = super::super::snapshot::PaneAgentSessionSnapshot {
             source: "herdr:pi".into(),
@@ -1177,7 +1311,13 @@ mod tests {
             resumed_sessions: &mut resumed,
         };
 
-        let startup = pane_restore_startup(Some(&session), Some(&history), &mut agent_restore);
+        let startup = pane_restore_startup(
+            Some(&session),
+            None,
+            std::path::Path::new("/a"),
+            Some(&history),
+            &mut agent_restore,
+        );
 
         assert!(startup.restore_plan.is_none());
         assert_eq!(startup.initial_history_ansi, Some("RESTORED_HISTORY\r\n"));
@@ -1332,6 +1472,7 @@ mod tests {
                                 kind: crate::agent_resume::AgentSessionRefKind::Id,
                                 value: "opencode-session".into(),
                             }),
+                            agent_resume: None,
                             launch_argv: None,
                         },
                     )]),
@@ -1413,6 +1554,7 @@ mod tests {
                                 agent_name: None,
                                 managed_agent_kind: None,
                                 agent_session: None,
+                                agent_resume: None,
                                 launch_argv: None,
                             },
                         ),
@@ -1424,6 +1566,7 @@ mod tests {
                                 agent_name: None,
                                 managed_agent_kind: None,
                                 agent_session: None,
+                                agent_resume: None,
                                 launch_argv: None,
                             },
                         ),
@@ -1477,6 +1620,7 @@ mod tests {
                     agent_name: None,
                     managed_agent_kind: None,
                     agent_session: None,
+                    agent_resume: None,
                     launch_argv: None,
                 },
             )
@@ -1492,6 +1636,7 @@ mod tests {
                 kind: crate::agent_resume::AgentSessionRefKind::Id,
                 value: "codex-session".into(),
             }),
+            agent_resume: None,
             launch_argv: None,
         };
         let snapshot = SessionSnapshot {
@@ -1643,6 +1788,7 @@ mod tests {
                                 kind: crate::agent_resume::AgentSessionRefKind::Id,
                                 value: "codex-session".into(),
                             }),
+                            agent_resume: None,
                             launch_argv: None,
                         },
                     )]),
@@ -1948,6 +2094,7 @@ mod tests {
                 agent_name: None,
                 managed_agent_kind: None,
                 agent_session: None,
+                agent_resume: None,
                 launch_argv: None,
             },
         );
diff --git a/src/persist/snapshot.rs b/src/persist/snapshot.rs
index d7db3e97..3518ef41 100644
--- a/src/persist/snapshot.rs
+++ b/src/persist/snapshot.rs
@@ -108,10 +108,19 @@ pub struct PaneSnapshot {
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub agent_session: Option<PaneAgentSessionSnapshot>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub agent_resume: Option<PaneAgentResumeSnapshot>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
     pub launch_argv: Option<Vec<String>>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct PaneAgentResumeSnapshot {
+    pub source: String,
+    pub agent: String,
+    pub argv: Vec<String>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct PaneAgentSessionSnapshot {
     pub source: String,
     pub agent: String,
@@ -362,6 +371,13 @@ fn capture_tab(
                     value: session.session_ref.value.clone(),
                 })
         });
+        let agent_resume = terminal
+            .and_then(|terminal| terminal.reported_resume())
+            .map(|resume| PaneAgentResumeSnapshot {
+                source: resume.source.clone(),
+                agent: resume.agent.clone(),
+                argv: resume.argv.clone(),
+            });
         panes.insert(
             id.raw(),
             PaneSnapshot {
@@ -370,6 +386,7 @@ fn capture_tab(
                 agent_name,
                 managed_agent_kind,
                 agent_session,
+                agent_resume,
                 launch_argv,
             },
         );
@@ -688,6 +705,7 @@ mod tests {
                 agent_name: None,
                 managed_agent_kind: None,
                 agent_session: None,
+                agent_resume: None,
                 launch_argv: None,
             },
         );
@@ -699,6 +717,7 @@ mod tests {
                 agent_name: None,
                 managed_agent_kind: None,
                 agent_session: None,
+                agent_resume: None,
                 launch_argv: None,
             },
         );
@@ -1287,6 +1306,40 @@ mod tests {
     }
 
     #[test]
+    fn capture_contract_includes_reported_agent_resume() {
+        let mut state = state_with_workspaces(&["one"]);
+        let root = state.workspaces[0].tabs[0].root_pane;
+        state.ensure_test_terminals();
+        let terminal_id = state.workspaces[0].tabs[0].panes[&root]
+            .attached_terminal_id
+            .clone();
+        let terminal = state.terminals.get_mut(&terminal_id).unwrap();
+        terminal.set_hook_authority(
+            "prime-agent".into(),
+            "prime-agent".into(),
+            crate::detect::AgentState::Idle,
+            None,
+            Some(1),
+        );
+        assert!(terminal.record_reported_resume(
+            "prime-agent",
+            "prime-agent",
+            Some(1),
+            vec!["prime-agent".into(), "--resume".into(), "a".into()],
+        ));
+
+        let snapshot = capture_from_state(&state);
+        let resume = snapshot.workspaces[0].tabs[0].panes[&root.raw()]
+            .agent_resume
+            .as_ref()
+            .expect("reported resume should be captured");
+
+        assert_eq!(resume.source, "prime-agent");
+        assert_eq!(resume.agent, "prime-agent");
+        assert_eq!(resume.argv, vec!["prime-agent", "--resume", "a"]);
+    }
+
+    #[test]
     fn capture_contract_preserves_restored_agent_session() {
         let mut state = state_with_workspaces(&["one"]);
         let root = state.workspaces[0].tabs[0].root_pane;
@@ -1350,6 +1403,7 @@ mod tests {
                 agent_name: None,
                 managed_agent_kind: None,
                 agent_session: None,
+                agent_resume: None,
                 launch_argv: None,
             },
         );
@@ -1363,6 +1417,7 @@ mod tests {
                 agent_name: None,
                 managed_agent_kind: None,
                 agent_session: None,
+                agent_resume: None,
                 launch_argv: None,
             },
         );