fix: allow foreground agent session takeover

refs #943

1 files changed, 205 insertions, 70 deletions

diff --git a/src/terminal/state.rs b/src/terminal/state.rs
index eebf0fb0..2b9ef4ad 100644
--- a/src/terminal/state.rs
+++ b/src/terminal/state.rs
@@ -410,9 +410,17 @@ impl TerminalState {
                 &agent_label,
                 &session_ref,
             );
-        if self.known_agent_label_conflicts_with_detected_agent(&agent_label)
-            || self.current_session_owner_conflicts(&source, &agent_label)
-        {
+        if self.known_agent_label_conflicts_with_detected_agent(&agent_label) {
+            return None;
+        }
+        let owner_conflicts = self.current_session_owner_conflicts(&source, &agent_label);
+        let foreground_takeover_allowed = owner_conflicts
+            && self.foreground_agent_confirms_hook_authority_takeover(
+                &source,
+                &agent_label,
+                &session_ref,
+            );
+        if owner_conflicts && !foreground_takeover_allowed {
             return None;
         }
         let session_ref = session_ref.map(|session_ref| {
@@ -446,6 +454,11 @@ impl TerminalState {
         let previous_state = self.state;
         let previous_presentation = self.effective_presentation_for_state_at(previous_state, now);
         let previous_session = self.current_session_identity_for_persistence();
+        if foreground_takeover_allowed {
+            self.suppress_current_full_lifecycle_hook_authority(
+                FullLifecycleHookSuppressionReason::HookClear,
+            );
+        }
         if session_ref.is_some() {
             if let Some(suppressed) = self.suppressed_full_lifecycle_hook_reports.remove(&source) {
                 if let Some(suppressed_ref) = suppressed.session_ref {
@@ -907,6 +920,13 @@ impl TerminalState {
         )
     }
 
+    fn session_start_source_is_recognized(session_start_source: Option<&str>) -> bool {
+        matches!(
+            session_start_source,
+            Some("startup" | "clear" | "resume" | "compact" | "new" | "fork")
+        )
+    }
+
     pub fn set_persisted_agent_session(
         &mut self,
         session: crate::agent_resume::PersistedAgentSession,
@@ -944,14 +964,15 @@ impl TerminalState {
             &agent_label,
             session_start_source.as_deref(),
         );
-        if self.current_session_owner_conflicts(&source, &agent_label)
-            && !(session_replacement_allowed
-                && self.foreground_agent_confirms_session_owner(
-                    &source,
-                    &agent_label,
-                    &session_ref,
-                ))
-        {
+        let owner_conflicts = self.current_session_owner_conflicts(&source, &agent_label);
+        let foreground_takeover_allowed = owner_conflicts
+            && self.foreground_agent_confirms_different_owner_takeover(
+                &source,
+                &agent_label,
+                &session_ref,
+                session_start_source.as_deref(),
+            );
+        if owner_conflicts && !foreground_takeover_allowed {
             return None;
         }
         if self
@@ -987,6 +1008,11 @@ impl TerminalState {
                 replaced_hook_session,
             );
             self.hook_authority = None;
+        } else if foreground_takeover_allowed {
+            self.suppress_current_full_lifecycle_hook_authority(
+                FullLifecycleHookSuppressionReason::HookClear,
+            );
+            self.hook_authority = None;
         }
         self.persisted_agent_session = Some(crate::agent_resume::PersistedAgentSession {
             source,
@@ -1014,6 +1040,28 @@ impl TerminalState {
             .is_some_and(|hook_agent| hook_agent != detected_agent)
     }
 
+    fn foreground_agent_confirms_different_owner_takeover(
+        &self,
+        source: &str,
+        agent_label: &str,
+        session_ref: &crate::agent_resume::AgentSessionRef,
+        session_start_source: Option<&str>,
+    ) -> bool {
+        Self::session_start_source_is_recognized(session_start_source)
+            && self.foreground_agent_confirms_session_owner(source, agent_label, session_ref)
+    }
+
+    fn foreground_agent_confirms_hook_authority_takeover(
+        &self,
+        source: &str,
+        agent_label: &str,
+        session_ref: &Option<crate::agent_resume::AgentSessionRef>,
+    ) -> bool {
+        session_ref.as_ref().is_some_and(|session_ref| {
+            self.foreground_agent_confirms_session_owner(source, agent_label, session_ref)
+        })
+    }
+
     fn foreground_agent_confirms_session_owner(
         &self,
         source: &str,
@@ -2435,7 +2483,7 @@ mod tests {
         let now = Instant::now();
         let mut terminal = test_terminal();
         terminal.set_detected_state(Some(Agent::Omp), AgentState::Idle);
-        terminal.set_hook_authority_with_session_ref(
+        terminal.set_hook_authority_with_custom_status_at(
             "herdr:omp".into(),
             "omp".into(),
             AgentState::Working,
@@ -2443,6 +2491,7 @@ mod tests {
             None,
             crate::agent_resume::AgentSessionRef::id("omp-old"),
             Some(1000),
+            now,
         );
         terminal.set_detected_state_with_screen_signals_at(
             Some(Agent::Omp),
@@ -3606,84 +3655,63 @@ mod tests {
 
     #[test]
     fn foreground_agent_session_replaces_stale_different_owner_session_ref() {
-        let mut terminal = test_terminal();
-        terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
-            source: "herdr:codex".into(),
-            agent: "codex".into(),
-            session_ref: crate::agent_resume::AgentSessionRef::id("codex-session").unwrap(),
-        });
-        terminal.set_detected_state(Some(Agent::Claude), AgentState::Idle);
+        for session_start_source in ["resume", "startup"] {
+            let mut terminal = test_terminal();
+            terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
+                source: "herdr:codex".into(),
+                agent: "codex".into(),
+                session_ref: crate::agent_resume::AgentSessionRef::id("codex-session").unwrap(),
+            });
+            terminal.set_detected_state(Some(Agent::Claude), AgentState::Idle);
 
-        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()),
-            )
-            .expect("foreground claude should replace stale codex session");
+            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(session_start_source.into()),
+                )
+                .unwrap_or_else(|| {
+                    panic!("{session_start_source} should replace stale codex session")
+                });
 
-        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"))
-        );
+            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")),
+                "{session_start_source} should store claude session"
+            );
+        }
     }
 
     #[test]
     fn foreground_agent_session_requires_lifecycle_source_to_replace_different_owner() {
-        let mut terminal = test_terminal();
-        terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
-            source: "herdr:codex".into(),
-            agent: "codex".into(),
-            session_ref: crate::agent_resume::AgentSessionRef::id("codex-session").unwrap(),
-        });
-        terminal.set_detected_state(Some(Agent::Claude), AgentState::Idle);
-
-        let mutation = terminal.set_agent_session_ref_for_session_start(
-            "herdr:claude".into(),
-            "claude".into(),
-            crate::agent_resume::AgentSessionRef::id("claude-session"),
-            Some(21),
-            None,
-        );
-
-        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:codex", "codex", "codex-session"))
-        );
-    }
-
-    #[test]
-    fn different_owner_session_ref_requires_matching_detected_agent() {
-        for detected_agent in [None, Some(Agent::Codex)] {
+        for session_start_source in [None, Some("other")] {
             let mut terminal = test_terminal();
             terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
                 source: "herdr:codex".into(),
                 agent: "codex".into(),
                 session_ref: crate::agent_resume::AgentSessionRef::id("codex-session").unwrap(),
             });
-            terminal.set_detected_state(detected_agent, AgentState::Idle);
+            terminal.set_detected_state(Some(Agent::Claude), AgentState::Idle);
 
             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()),
+                session_start_source.map(str::to_string),
             );
 
-            assert!(mutation.is_none(), "{detected_agent:?} should not replace");
+            assert!(
+                mutation.is_none(),
+                "{session_start_source:?} should not replace"
+            );
             assert_eq!(
                 terminal.persisted_agent_session.as_ref().map(|session| (
                     session.source.as_str(),
@@ -3696,6 +3724,42 @@ mod tests {
     }
 
     #[test]
+    fn different_owner_session_ref_requires_matching_detected_agent() {
+        for session_start_source in ["startup", "resume"] {
+            for detected_agent in [None, Some(Agent::Codex)] {
+                let mut terminal = test_terminal();
+                terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
+                    source: "herdr:codex".into(),
+                    agent: "codex".into(),
+                    session_ref: crate::agent_resume::AgentSessionRef::id("codex-session").unwrap(),
+                });
+                terminal.set_detected_state(detected_agent, AgentState::Idle);
+
+                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(session_start_source.into()),
+                );
+
+                assert!(
+                    mutation.is_none(),
+                    "{session_start_source} with {detected_agent:?} should not replace"
+                );
+                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:codex", "codex", "codex-session"))
+                );
+            }
+        }
+    }
+
+    #[test]
     fn custom_session_report_does_not_replace_different_owner_session_ref() {
         let mut terminal = test_terminal();
         terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
@@ -3725,6 +3789,77 @@ mod tests {
     }
 
     #[test]
+    fn foreground_agent_session_replaces_stale_different_owner_hook_authority() {
+        let mut terminal = test_terminal();
+        let now = std::time::Instant::now();
+        terminal
+            .set_hook_authority_with_custom_status_at(
+                "herdr:opencode".into(),
+                "opencode".into(),
+                AgentState::Working,
+                None,
+                None,
+                crate::agent_resume::AgentSessionRef::id("opencode-session"),
+                Some(20),
+                now + Duration::from_millis(1),
+            )
+            .expect("initial hook authority should be accepted");
+        terminal.set_detected_state_with_screen_signals_at(
+            Some(Agent::Codex),
+            AgentState::Idle,
+            false,
+            false,
+            false,
+            false,
+            now,
+        );
+
+        let mutation = terminal
+            .set_agent_session_ref_for_session_start(
+                "herdr:codex".into(),
+                "codex".into(),
+                crate::agent_resume::AgentSessionRef::id("codex-session"),
+                Some(21),
+                Some("startup".into()),
+            )
+            .expect("foreground codex should replace stale hook authority");
+
+        assert!(mutation.session_ref_changed);
+        assert!(terminal.hook_authority.is_none());
+        assert_eq!(
+            terminal.current_session_identity_for_persistence(),
+            Some((
+                "herdr:codex".into(),
+                "codex".into(),
+                crate::agent_resume::AgentSessionRefKind::Id,
+                "codex-session".into()
+            ))
+        );
+        let late_old_session = terminal.set_hook_authority_with_session_ref(
+            "herdr:opencode".into(),
+            "opencode".into(),
+            AgentState::Working,
+            None,
+            None,
+            crate::agent_resume::AgentSessionRef::id("opencode-session"),
+            Some(22),
+        );
+        assert!(late_old_session.is_none());
+
+        terminal.set_detected_state(Some(Agent::OpenCode), AgentState::Idle);
+        let fresh_session = terminal.set_hook_authority_with_session_ref(
+            "herdr:opencode".into(),
+            "opencode".into(),
+            AgentState::Working,
+            None,
+            None,
+            crate::agent_resume::AgentSessionRef::id("opencode-new-session"),
+            Some(23),
+        );
+        assert!(fresh_session.is_some());
+    }
+
+    #[test]
     fn different_owner_full_lifecycle_hook_does_not_replace_existing_session_ref() {
         let mut terminal = test_terminal();
         terminal