fix: protect root agent restore sessions

refs #712

1 files changed, 273 insertions, 76 deletions

diff --git a/src/terminal/state.rs b/src/terminal/state.rs
index 01b6bdfe..b5a05168 100644
--- a/src/terminal/state.rs
+++ b/src/terminal/state.rs
@@ -264,6 +264,14 @@ impl TerminalState {
             }
             self.hook_authority = None;
         }
+        if process_exited
+            && self
+                .persisted_agent_session
+                .as_ref()
+                .is_some_and(|session| crate::detect::parse_agent_label(&session.agent) == agent)
+        {
+            self.persisted_agent_session = None;
+        }
         if self.hook_authority_not_newer_than(now)
             && (self.hook_authority_conflicts_with_detected_agent(agent)
                 || (previous_detected_agent.is_some()
@@ -273,19 +281,20 @@ impl TerminalState {
                             == previous_detected_agent
                     })))
         {
+            let durable_session = self.hook_authority.as_ref().and_then(|authority| {
+                authority.session_ref.as_ref().map(|session_ref| {
+                    crate::agent_resume::PersistedAgentSession {
+                        source: authority.source.clone(),
+                        agent: authority.agent_label.clone(),
+                        session_ref: session_ref.clone(),
+                    }
+                })
+            });
             self.suppress_current_full_lifecycle_hook_authority(
                 FullLifecycleHookSuppressionReason::HookClear,
             );
             self.hook_authority = None;
-        }
-        let detected_agent_changed_or_disappeared =
-            previous_detected_agent.is_some() && agent != previous_detected_agent;
-        let persisted_agent_was_previously_detected =
-            self.persisted_agent_session_belongs_to_detected_agent(previous_detected_agent);
-        if self.persisted_agent_session_conflicts_with_detected_agent(agent)
-            || detected_agent_changed_or_disappeared && persisted_agent_was_previously_detected
-        {
-            self.persisted_agent_session = None;
+            self.persisted_agent_session = durable_session;
         }
         TerminalStateMutation {
             effective_state_change: self.recompute_effective_state(
@@ -388,11 +397,13 @@ impl TerminalState {
                 &agent_label,
                 &session_ref,
             );
-        if self.known_agent_label_conflicts_with_detected_agent(&agent_label) {
+        if self.known_agent_label_conflicts_with_detected_agent(&agent_label)
+            || self.current_session_owner_conflicts(&source, &agent_label)
+        {
             return None;
         }
         let session_ref = session_ref.map(|session_ref| {
-            self.conflicting_current_session_ref(&source, &agent_label, &session_ref, None)
+            self.conflicting_same_owner_session_ref(&source, &agent_label, &session_ref, None)
                 .unwrap_or(session_ref)
         });
         if self.live_full_lifecycle_hook_authority_conflicts_with_session(
@@ -481,32 +492,6 @@ impl TerminalState {
             && !self.hook_authority_conflicts_with_detected_agent(detected_agent)
     }
 
-    fn persisted_agent_session_conflicts_with_detected_agent(
-        &self,
-        detected_agent: Option<Agent>,
-    ) -> bool {
-        let Some(detected_agent) = detected_agent else {
-            return false;
-        };
-        self.persisted_agent_session
-            .as_ref()
-            .and_then(|session| crate::detect::parse_agent_label(&session.agent))
-            .is_some_and(|agent| agent != detected_agent)
-    }
-
-    fn persisted_agent_session_belongs_to_detected_agent(
-        &self,
-        detected_agent: Option<Agent>,
-    ) -> bool {
-        let Some(detected_agent) = detected_agent else {
-            return false;
-        };
-        self.persisted_agent_session
-            .as_ref()
-            .and_then(|session| crate::detect::parse_agent_label(&session.agent))
-            .is_some_and(|agent| agent == detected_agent)
-    }
-
     fn persisted_agent_session_matches(&self, source: &str, agent: &str) -> bool {
         self.persisted_agent_session
             .as_ref()
@@ -802,7 +787,15 @@ impl TerminalState {
         })
     }
 
-    fn conflicting_current_session_ref(
+    fn current_session_owner_conflicts(&self, source: &str, agent_label: &str) -> bool {
+        self.current_session_identity_for_persistence().is_some_and(
+            |(current_source, current_agent, _, _)| {
+                current_source != source || current_agent != agent_label
+            },
+        )
+    }
+
+    fn conflicting_same_owner_session_ref(
         &self,
         source: &str,
         agent_label: &str,
@@ -815,7 +808,7 @@ impl TerminalState {
                     && current_agent == agent_label
                     && current_kind == crate::agent_resume::AgentSessionRefKind::Id
                     && session_ref.kind == crate::agent_resume::AgentSessionRefKind::Id
-                    && (current_kind != session_ref.kind || current_value != session_ref.value)
+                    && current_value != session_ref.value
                     && !Self::session_start_source_allows_session_replacement(
                         source,
                         agent_label,
@@ -834,9 +827,18 @@ impl TerminalState {
         agent_label: &str,
         session_start_source: Option<&str>,
     ) -> bool {
-        source == "herdr:claude"
-            && agent_label == "claude"
-            && matches!(session_start_source, Some("clear" | "resume" | "compact"))
+        matches!(
+            (source, agent_label, session_start_source),
+            (
+                "herdr:claude",
+                "claude",
+                Some("clear" | "resume" | "compact")
+            ) | (
+                "herdr:codex",
+                "codex",
+                Some("startup" | "clear" | "resume" | "compact")
+            )
+        )
     }
 
     pub fn set_persisted_agent_session(
@@ -871,14 +873,15 @@ impl TerminalState {
         if self.known_agent_label_conflicts_with_detected_agent(&agent_label) {
             return None;
         }
-        if self
-            .conflicting_current_session_ref(
-                &source,
-                &agent_label,
-                &session_ref,
-                session_start_source.as_deref(),
-            )
-            .is_some()
+        if self.current_session_owner_conflicts(&source, &agent_label)
+            || self
+                .conflicting_same_owner_session_ref(
+                    &source,
+                    &agent_label,
+                    &session_ref,
+                    session_start_source.as_deref(),
+                )
+                .is_some()
         {
             return None;
         }
@@ -1009,6 +1012,10 @@ impl TerminalState {
         if !matches_current_agent && !matches_persisted_session {
             return None;
         }
+        let preserve_foreign_persisted_session = self
+            .persisted_agent_session
+            .as_ref()
+            .is_some_and(|session| session.source != source || session.agent != agent_label);
 
         let now = Instant::now();
         let previous_agent_label = self.effective_agent_label().map(str::to_string);
@@ -1026,7 +1033,10 @@ impl TerminalState {
         self.fallback_visible_blocker = false;
         self.fallback_observed_at = None;
         self.hook_authority = None;
-        self.persisted_agent_session = None;
+        if !preserve_foreign_persisted_session {
+            self.persisted_agent_session = None;
+        }
+        let current_session = self.current_session_identity_for_persistence();
         Some(TerminalStateMutation {
             effective_state_change: self.recompute_effective_state(
                 previous_agent_label,
@@ -1035,7 +1045,7 @@ impl TerminalState {
                 previous_presentation,
                 now,
             ),
-            session_ref_changed: previous_session.is_some(),
+            session_ref_changed: previous_session != current_session,
         })
     }
 
@@ -3097,22 +3107,23 @@ mod tests {
     }
 
     #[test]
-    fn accepted_hook_report_marks_changed_when_session_identity_changes() {
+    fn accepted_hook_report_marks_changed_when_same_owner_session_identity_changes() {
         let mut terminal = test_terminal();
         terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
-            source: "herdr:opencode".into(),
-            agent: "opencode".into(),
-            session_ref: crate::agent_resume::AgentSessionRef::id("same-session").unwrap(),
+            source: "herdr:pi".into(),
+            agent: "pi".into(),
+            session_ref: crate::agent_resume::AgentSessionRef::path(test_session_path("old.jsonl"))
+                .unwrap(),
         });
 
         let mutation = terminal
             .set_hook_authority_with_session_ref(
-                "herdr:hermes".into(),
-                "hermes".into(),
+                "herdr:pi".into(),
+                "pi".into(),
                 AgentState::Working,
                 None,
                 None,
-                crate::agent_resume::AgentSessionRef::id("same-session"),
+                crate::agent_resume::AgentSessionRef::path(test_session_path("new.jsonl")),
                 Some(20),
             )
             .expect("accepted report");
@@ -3223,6 +3234,106 @@ mod tests {
     }
 
     #[test]
+    fn codex_lifecycle_session_ref_replaces_existing_session_ref() {
+        for session_start_source in ["startup", "clear", "resume", "compact"] {
+            let mut terminal = test_terminal();
+            terminal
+                .set_agent_session_ref(
+                    "herdr:codex".into(),
+                    "codex".into(),
+                    crate::agent_resume::AgentSessionRef::id("codex-session"),
+                    Some(20),
+                )
+                .expect("initial session should be accepted");
+
+            let next_session = format!("codex-{session_start_source}-session");
+            let mutation = terminal
+                .set_agent_session_ref_for_session_start(
+                    "herdr:codex".into(),
+                    "codex".into(),
+                    crate::agent_resume::AgentSessionRef::id(&next_session),
+                    Some(21),
+                    Some(session_start_source.into()),
+                )
+                .unwrap_or_else(|| panic!("{session_start_source} should replace the session"));
+
+            assert!(mutation.session_ref_changed);
+            assert_eq!(
+                terminal
+                    .persisted_agent_session
+                    .as_ref()
+                    .map(|session| session.session_ref.value.as_str()),
+                Some(next_session.as_str())
+            );
+        }
+    }
+
+    #[test]
+    fn different_owner_session_ref_does_not_replace_existing_session_ref() {
+        let mut terminal = test_terminal();
+        terminal
+            .set_agent_session_ref(
+                "herdr:droid".into(),
+                "droid".into(),
+                crate::agent_resume::AgentSessionRef::id("droid-session"),
+                Some(20),
+            )
+            .expect("initial session should be accepted");
+
+        let mutation = terminal.set_agent_session_ref_for_session_start(
+            "herdr:claude".into(),
+            "claude".into(),
+            crate::agent_resume::AgentSessionRef::id("claude-session"),
+            Some(21),
+            Some("resume".into()),
+        );
+
+        assert!(mutation.is_none());
+        assert_eq!(
+            terminal.persisted_agent_session.as_ref().map(|session| (
+                session.source.as_str(),
+                session.agent.as_str(),
+                session.session_ref.value.as_str()
+            )),
+            Some(("herdr:droid", "droid", "droid-session"))
+        );
+    }
+
+    #[test]
+    fn different_owner_full_lifecycle_hook_does_not_replace_existing_session_ref() {
+        let mut terminal = test_terminal();
+        terminal
+            .set_agent_session_ref(
+                "herdr:droid".into(),
+                "droid".into(),
+                crate::agent_resume::AgentSessionRef::id("droid-session"),
+                Some(20),
+            )
+            .expect("initial session should be accepted");
+
+        let mutation = terminal.set_hook_authority_with_session_ref(
+            "herdr:pi".into(),
+            "pi".into(),
+            AgentState::Working,
+            None,
+            None,
+            crate::agent_resume::AgentSessionRef::path("/tmp/pi-session.jsonl"),
+            Some(21),
+        );
+
+        assert!(mutation.is_none());
+        assert!(terminal.hook_authority.is_none());
+        assert_eq!(
+            terminal.persisted_agent_session.as_ref().map(|session| (
+                session.source.as_str(),
+                session.agent.as_str(),
+                session.session_ref.value.as_str()
+            )),
+            Some(("herdr:droid", "droid", "droid-session"))
+        );
+    }
+
+    #[test]
     fn repeated_same_agent_session_ref_is_accepted_without_session_change() {
         let mut terminal = test_terminal();
         terminal
@@ -3286,7 +3397,7 @@ mod tests {
     }
 
     #[test]
-    fn different_same_agent_session_ref_is_accepted_after_detection_clears_current_session() {
+    fn detected_agent_clear_does_not_clear_current_session_ref() {
         let mut terminal = test_terminal();
         terminal.set_detected_state(Some(Agent::Claude), AgentState::Working);
         terminal
@@ -3299,24 +3410,22 @@ mod tests {
             .expect("initial session should be accepted");
 
         let clear = terminal.set_detected_state_with_mutation(None, AgentState::Unknown);
-        assert!(clear.session_ref_changed);
+        assert!(!clear.session_ref_changed);
 
-        let mutation = terminal
-            .set_agent_session_ref(
-                "herdr:claude".into(),
-                "claude".into(),
-                crate::agent_resume::AgentSessionRef::id("new-session"),
-                Some(21),
-            )
-            .expect("new session should be accepted after clear");
+        let mutation = terminal.set_agent_session_ref(
+            "herdr:claude".into(),
+            "claude".into(),
+            crate::agent_resume::AgentSessionRef::id("new-session"),
+            Some(21),
+        );
 
-        assert!(mutation.session_ref_changed);
+        assert!(mutation.is_none());
         assert_eq!(
             terminal
                 .persisted_agent_session
                 .as_ref()
                 .map(|session| session.session_ref.value.as_str()),
-            Some("new-session")
+            Some("claude-session")
         );
     }
 
@@ -3383,6 +3492,86 @@ mod tests {
     }
 
     #[test]
+    fn release_agent_preserves_foreign_persisted_session_ref() {
+        let mut terminal = test_terminal();
+        terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
+            source: "herdr:claude".into(),
+            agent: "claude".into(),
+            session_ref: crate::agent_resume::AgentSessionRef::id("claude-session").unwrap(),
+        });
+        terminal.set_detected_state(Some(Agent::Pi), AgentState::Idle);
+
+        let mutation = terminal
+            .release_agent_with_mutation("herdr:pi", "pi", Some(21))
+            .expect("visible agent release should be accepted");
+
+        assert!(!mutation.session_ref_changed);
+        assert_eq!(
+            terminal.persisted_agent_session.as_ref().map(|session| (
+                session.source.as_str(),
+                session.agent.as_str(),
+                session.session_ref.value.as_str()
+            )),
+            Some(("herdr:claude", "claude", "claude-session"))
+        );
+    }
+
+    #[test]
+    fn process_exit_clears_matching_persisted_session_ref() {
+        let mut terminal = test_terminal();
+        terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
+            source: "herdr:pi".into(),
+            agent: "pi".into(),
+            session_ref: crate::agent_resume::AgentSessionRef::path(test_session_path("pi.jsonl"))
+                .unwrap(),
+        });
+        terminal.set_detected_state(Some(Agent::Pi), AgentState::Working);
+
+        let mutation = terminal.set_detected_state_with_screen_signals_at(
+            Some(Agent::Pi),
+            AgentState::Idle,
+            false,
+            false,
+            false,
+            true,
+            std::time::Instant::now(),
+        );
+
+        assert!(mutation.session_ref_changed);
+        assert!(terminal.persisted_agent_session.is_none());
+    }
+
+    #[test]
+    fn process_exit_preserves_foreign_persisted_session_ref() {
+        let mut terminal = test_terminal();
+        terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
+            source: "herdr:claude".into(),
+            agent: "claude".into(),
+            session_ref: crate::agent_resume::AgentSessionRef::id("claude-session").unwrap(),
+        });
+        terminal.set_detected_state(Some(Agent::Pi), AgentState::Working);
+
+        let mutation = terminal.set_detected_state_with_screen_signals_at(
+            Some(Agent::Pi),
+            AgentState::Idle,
+            false,
+            false,
+            false,
+            true,
+            std::time::Instant::now(),
+        );
+
+        assert!(!mutation.session_ref_changed);
+        assert_eq!(
+            terminal
+                .persisted_agent_session
+                .as_ref()
+                .map(|session| session.session_ref.value.as_str()),
+            Some("claude-session")
+        );
+    }
+
+    #[test]
     fn respawn_cleanup_resets_restored_agent_status() {
         let mut terminal = test_terminal();
         terminal.respawn_shell_on_exit = true;
@@ -3404,7 +3593,7 @@ mod tests {
     }
 
     #[test]
-    fn detected_conflict_clears_session_ref() {
+    fn detected_conflict_clears_live_hook_but_preserves_session_ref() {
         let mut terminal = test_terminal();
         terminal.set_hook_authority_with_session_ref(
             "herdr:claude".into(),
@@ -3419,8 +3608,16 @@ mod tests {
         let mutation =
             terminal.set_detected_state_with_mutation(Some(Agent::Grok), AgentState::Idle);
 
-        assert!(mutation.session_ref_changed);
+        assert!(!mutation.session_ref_changed);
         assert!(terminal.hook_authority.is_none());
+        assert_eq!(
+            terminal.persisted_agent_session.as_ref().map(|session| (
+                session.source.as_str(),
+                session.agent.as_str(),
+                session.session_ref.value.as_str()
+            )),
+            Some(("herdr:claude", "claude", "claude-session"))
+        );
     }
 
     #[test]
@@ -3446,7 +3643,7 @@ mod tests {
     }
 
     #[test]
-    fn detected_agent_disappearance_clears_matching_persisted_session_ref() {
+    fn detected_agent_disappearance_preserves_matching_persisted_session_ref() {
         let mut terminal = test_terminal();
         terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
             source: "herdr:opencode".into(),
@@ -3460,8 +3657,8 @@ mod tests {
         assert!(terminal.persisted_agent_session.is_some());
 
         let second = terminal.set_detected_state_with_mutation(None, AgentState::Unknown);
-        assert!(second.session_ref_changed);
-        assert!(terminal.persisted_agent_session.is_none());
+        assert!(!second.session_ref_changed);
+        assert!(terminal.persisted_agent_session.is_some());
     }
 
     #[test]