fix: recover omp resume and blocked states

refs #879

4 files changed, 356 insertions, 29 deletions

diff --git a/src/agent_resume.rs b/src/agent_resume.rs
index a157f06a..5c1001cb 100644
--- a/src/agent_resume.rs
+++ b/src/agent_resume.rs
@@ -71,7 +71,7 @@ pub fn session_ref_from_report(
 
 pub fn normalize_session_start_source(value: Option<String>) -> Option<String> {
     match value.as_deref().map(str::trim) {
-        Some(source @ ("startup" | "resume" | "clear" | "compact" | "new")) => {
+        Some(source @ ("startup" | "resume" | "clear" | "compact" | "new" | "fork")) => {
             Some(source.to_string())
         }
         _ => None,
@@ -525,6 +525,10 @@ mod tests {
             Some("new".into())
         );
         assert_eq!(
+            normalize_session_start_source(Some("fork".into())),
+            Some("fork".into())
+        );
+        assert_eq!(
             normalize_session_start_source(Some(" resume ".into())),
             Some("resume".into())
         );
diff --git a/src/integration/assets/omp/herdr-agent-state.ts b/src/integration/assets/omp/herdr-agent-state.ts
index e0cfab5e..21c218b7 100644
--- a/src/integration/assets/omp/herdr-agent-state.ts
+++ b/src/integration/assets/omp/herdr-agent-state.ts
@@ -16,7 +16,9 @@ function enabled() {
   return HERDR_ENV === "1" && !!socketPath && !!paneId;
 }
 
-function sendRequest(request: unknown): Promise<void> {
+let requestQueue = Promise.resolve();
+
+function sendRequestNow(request: unknown): Promise<void> {
   if (!enabled()) {
     return Promise.resolve();
   }
@@ -40,6 +42,14 @@ function sendRequest(request: unknown): Promise<void> {
   });
 }
 
+function sendRequest(request: unknown): Promise<void> {
+  requestQueue = requestQueue.then(
+    () => sendRequestNow(request),
+    () => sendRequestNow(request),
+  );
+  return requestQueue;
+}
+
 type AgentState = "working" | "blocked" | "idle";
 
 type QueuedState = {
@@ -110,7 +120,7 @@ function currentSessionRef(): Record<string, unknown> | undefined {
   return undefined;
 }
 
-function reportSession(): Promise<void> {
+function reportSession(sessionStartSource = "startup"): Promise<void> {
   const sessionRef = currentSessionRef();
   if (!sessionRef) {
     return Promise.resolve();
@@ -124,6 +134,7 @@ function reportSession(): Promise<void> {
       source,
       agent: "omp",
       seq: nextReportSeq(),
+      session_start_source: sessionStartSource,
       ...sessionRef,
     },
   });
@@ -221,6 +232,15 @@ function retryableErrorMessage(event: any): string | undefined {
   return errorMessage || "retryable provider error";
 }
 
+function askBlockedMessage(args: any): string {
+  const questions = Array.isArray(args?.questions) ? args.questions : [];
+  const firstQuestion = questions.find((question: any) => typeof question?.question === "string");
+  if (firstQuestion?.question) {
+    return firstQuestion.question;
+  }
+  return "waiting for user input";
+}
+
 export default function (pi) {
   if (!enabled()) {
     return;
@@ -306,37 +326,68 @@ export default function (pi) {
     retryTimer.unref?.();
   }
 
+  function activateRootSession(ctx: any, sessionStartSource = "startup"): boolean {
+    if (ctx?.hasUI !== true) {
+      return false;
+    }
+    rootSession = true;
+    updateSessionRef(ctx);
+    void reportSession(sessionStartSource);
+    return true;
+  }
+
+  function resetSessionState() {
+    clearPendingTimers();
+    clearFailureState();
+    agentActive = false;
+    blockedCount = 0;
+    blockedMessage = undefined;
+  }
+
+  function activateBlocked(message: string | undefined) {
+    clearPendingTimers();
+    blockedCount += 1;
+    blockedMessage = message;
+    publishState();
+  }
+
+  function deactivateBlocked() {
+    blockedCount = Math.max(0, blockedCount - 1);
+    if (blockedCount === 0) {
+      blockedMessage = undefined;
+    }
+    publishState();
+  }
+
   pi.events.on("herdr:blocked", (data) => {
     if (!rootSession) {
       return;
     }
     if (!data?.active) {
-      blockedCount = Math.max(0, blockedCount - 1);
-      if (blockedCount === 0) {
-        blockedMessage = undefined;
-      }
-      publishState();
+      deactivateBlocked();
       return;
     }
 
-    clearPendingTimers();
-    blockedCount += 1;
-    blockedMessage = data.label;
-    publishState();
+    activateBlocked(data.label);
   });
 
   pi.on("session_start", (_event, ctx) => {
-    if (ctx?.hasUI !== true) {
+    if (!activateRootSession(ctx)) {
       return;
     }
-    rootSession = true;
-    updateSessionRef(ctx);
-    void reportSession();
+    publishState(true);
+  });
+
+  pi.on("session_switch", (event, ctx) => {
+    if (!activateRootSession(ctx, event?.reason || "resume")) {
+      return;
+    }
+    resetSessionState();
     publishState(true);
   });
 
   pi.on("agent_start", (_event, ctx) => {
-    if (!rootSession) {
+    if (!rootSession && !activateRootSession(ctx)) {
       return;
     }
     updateSessionRef(ctx);
@@ -347,6 +398,41 @@ export default function (pi) {
     publishState();
   });
 
+  pi.on("tool_approval_requested", (event, ctx) => {
+    if (!rootSession && !activateRootSession(ctx)) {
+      return;
+    }
+    const label = event?.reason || `${event?.toolName || "Tool"} approval`;
+    activateBlocked(label);
+  });
+
+  pi.on("tool_approval_resolved", (_event, ctx) => {
+    if (!rootSession && !activateRootSession(ctx)) {
+      return;
+    }
+    deactivateBlocked();
+  });
+
+  pi.on("tool_execution_start", (event, ctx) => {
+    if (event?.toolName !== "ask") {
+      return;
+    }
+    if (!rootSession && !activateRootSession(ctx)) {
+      return;
+    }
+    activateBlocked(askBlockedMessage(event.args));
+  });
+
+  pi.on("tool_execution_end", (event, ctx) => {
+    if (event?.toolName !== "ask") {
+      return;
+    }
+    if (!rootSession && !activateRootSession(ctx)) {
+      return;
+    }
+    deactivateBlocked();
+  });
+
   pi.on("agent_end", (event) => {
     if (!rootSession) {
       return;
diff --git a/src/integration/tests.rs b/src/integration/tests.rs
index 4c40e2af..2a8a7be3 100644
--- a/src/integration/tests.rs
+++ b/src/integration/tests.rs
@@ -2671,20 +2671,134 @@ fn omp_extension_refreshes_session_ref_before_agent_start_state() {
     assert!(report_session < publish_state);
 }
 
-#[test]
-fn omp_session_hook_ignores_non_ui_sessions() {
-    let session_start_handler = OMP_EXTENSION_ASSET
-        .find("pi.on(\"session_start\"")
-        .expect("omp extension registers session_start handler");
-    let non_ui_guard = OMP_EXTENSION_ASSET
+fn omp_handler(event: &str) -> &'static str {
+    let start = OMP_EXTENSION_ASSET
+        .find(&format!("pi.on(\"{event}\""))
+        .unwrap_or_else(|| panic!("omp extension registers {event} handler"));
+    let rest = &OMP_EXTENSION_ASSET[start..];
+    let end = rest[1..]
+        .find("\n\n  pi.")
+        .map(|offset| offset + 1)
+        .unwrap_or(rest.len());
+    &rest[..end]
+}
+
+#[test]
+fn omp_root_activation_requires_ui_context() {
+    let activator = OMP_EXTENSION_ASSET
+        .find("function activateRootSession(ctx: any, sessionStartSource = \"startup\"): boolean")
+        .expect("omp extension should centralize root session activation");
+    let helper = &OMP_EXTENSION_ASSET[activator..];
+    let non_ui_guard = helper
         .find("ctx?.hasUI !== true")
-        .expect("omp extension checks UI context");
-    let session_report = OMP_EXTENSION_ASSET
-        .find("void reportSession()")
+        .expect("omp extension checks UI context before activating");
+    let root_session = helper
+        .find("rootSession = true;")
+        .expect("omp extension activates root session after UI guard");
+    let session_report = helper
+        .find("void reportSession(sessionStartSource);")
         .expect("omp extension reports root session");
 
-    assert!(session_start_handler < non_ui_guard);
-    assert!(non_ui_guard < session_report);
+    assert!(non_ui_guard < root_session);
+    assert!(root_session < session_report);
+}
+
+#[test]
+fn omp_session_start_and_switch_use_root_activation() {
+    let session_start = OMP_EXTENSION_ASSET
+        .find("pi.on(\"session_start\", (_event, ctx)")
+        .expect("omp extension registers session_start handler");
+    let session_start_handler = &OMP_EXTENSION_ASSET[session_start..];
+    session_start_handler
+        .find("if (!activateRootSession(ctx))")
+        .expect("omp session_start handler should activate root session");
+
+    let session_switch = OMP_EXTENSION_ASSET
+        .find("pi.on(\"session_switch\", (event, ctx)")
+        .expect("omp extension registers session_switch handler");
+    let session_switch_handler = &OMP_EXTENSION_ASSET[session_switch..];
+    session_switch_handler
+        .find("if (!activateRootSession(ctx, event?.reason || \"resume\"))")
+        .expect("omp session_switch handler should activate root session with switch reason");
+}
+
+#[test]
+fn omp_session_reports_include_start_source() {
+    let report_session = OMP_EXTENSION_ASSET
+        .find("function reportSession(sessionStartSource = \"startup\"): Promise<void>")
+        .expect("omp extension should label session reports with a lifecycle source");
+    let helper = &OMP_EXTENSION_ASSET[report_session..];
+    let session_source = helper
+        .find("session_start_source: sessionStartSource")
+        .expect("omp session reports should include the lifecycle source");
+    let session_ref = helper
+        .find("...sessionRef")
+        .expect("omp session reports should include the native session ref");
+
+    assert!(session_source < session_ref);
+}
+
+#[test]
+fn omp_socket_requests_are_serialized() {
+    let queue = OMP_EXTENSION_ASSET
+        .find("let requestQueue = Promise.resolve();")
+        .expect("omp extension should keep socket reports ordered");
+    let send_request = OMP_EXTENSION_ASSET[queue..]
+        .find("function sendRequest(request: unknown): Promise<void>")
+        .expect("omp extension should wrap socket sends in an ordered queue");
+    let queued_send = OMP_EXTENSION_ASSET[queue + send_request..]
+        .find("requestQueue = requestQueue.then(")
+        .expect("omp extension should serialize socket requests through the queue");
+    let raw_send = OMP_EXTENSION_ASSET[queue + send_request..]
+        .find("sendRequestNow(request)")
+        .expect("omp extension should enqueue the raw socket send");
+
+    assert!(queued_send < raw_send);
+}
+
+#[test]
+fn omp_runtime_events_can_activate_root_session_after_resume() {
+    for event in [
+        "agent_start",
+        "tool_approval_requested",
+        "tool_approval_resolved",
+        "tool_execution_start",
+        "tool_execution_end",
+    ] {
+        let handler = omp_handler(event);
+        handler
+            .find("!rootSession && !activateRootSession(ctx)")
+            .unwrap_or_else(|| panic!("omp {event} handler should recover missing root session"));
+    }
+}
+
+#[test]
+fn omp_ask_and_approval_events_report_blocked_state() {
+    let approval_handler = omp_handler("tool_approval_requested");
+    approval_handler
+        .find("activateBlocked(label);")
+        .expect("approval requests should block the pane");
+
+    let approval_resolved = omp_handler("tool_approval_resolved");
+    approval_resolved
+        .find("deactivateBlocked();")
+        .expect("approval resolution should unblock the pane");
+
+    let ask_handler = omp_handler("tool_execution_start");
+    ask_handler
+        .find("event?.toolName !== \"ask\"")
+        .expect("tool execution handler should only treat Ask as blocked");
+    ask_handler
+        .find("activateBlocked(askBlockedMessage(event.args));")
+        .expect("Ask start should block the pane");
+
+    let ask_end_handler = omp_handler("tool_execution_end");
+    ask_end_handler
+        .find("event?.toolName !== \"ask\"")
+        .expect("tool execution end should only treat Ask as blocked");
+    ask_end_handler
+        .find("deactivateBlocked();")
+        .expect("Ask end should unblock the pane");
 }
 
 #[test]
diff --git a/src/terminal/state.rs b/src/terminal/state.rs
index 8fea0542..3e536e6e 100644
--- a/src/terminal/state.rs
+++ b/src/terminal/state.rs
@@ -692,6 +692,26 @@ impl TerminalState {
             .is_some_and(|(current, incoming)| current != incoming)
     }
 
+    fn same_owner_full_lifecycle_hook_authority_session_ref(
+        &self,
+        source: &str,
+        agent_label: &str,
+        session_ref: &crate::agent_resume::AgentSessionRef,
+    ) -> Option<crate::agent_resume::AgentSessionRef> {
+        let authority = self.hook_authority.as_ref()?;
+        if !crate::detect::full_lifecycle_hook_authority(&authority.source, &authority.agent_label)
+            || authority.source != source
+            || authority.agent_label != agent_label
+        {
+            return None;
+        }
+        authority
+            .session_ref
+            .as_ref()
+            .filter(|current| *current != session_ref)
+            .cloned()
+    }
+
     fn clear_full_lifecycle_hook_suppression_for_detected_agent(
         &mut self,
         previous_detected_agent: Option<Agent>,
@@ -851,6 +871,11 @@ impl TerminalState {
                 "codex",
                 Some("startup" | "clear" | "resume" | "compact")
             ) | ("herdr:opencode", "opencode", Some("new"))
+                | (
+                    "herdr:omp",
+                    "omp",
+                    Some("startup" | "new" | "resume" | "fork")
+                )
         )
     }
 
@@ -886,6 +911,11 @@ impl TerminalState {
         if self.known_agent_label_conflicts_with_detected_agent(&agent_label) {
             return None;
         }
+        let session_replacement_allowed = Self::session_start_source_allows_session_replacement(
+            &source,
+            &agent_label,
+            session_start_source.as_deref(),
+        );
         if self.current_session_owner_conflicts(&source, &agent_label)
             || self
                 .conflicting_same_owner_session_ref(
@@ -898,8 +928,29 @@ impl TerminalState {
         {
             return None;
         }
+        let replaced_hook_session = self.same_owner_full_lifecycle_hook_authority_session_ref(
+            &source,
+            &agent_label,
+            &session_ref,
+        );
+        if replaced_hook_session.is_some() && !session_replacement_allowed {
+            return None;
+        }
 
+        let now = Instant::now();
+        let previous_agent_label = self.effective_agent_label().map(str::to_string);
+        let previous_known_agent = self.effective_known_agent();
+        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 let Some(replaced_hook_session) = replaced_hook_session {
+            self.remember_stale_full_lifecycle_hook_session(
+                source.clone(),
+                agent_label.clone(),
+                replaced_hook_session,
+            );
+            self.hook_authority = None;
+        }
         self.persisted_agent_session = Some(crate::agent_resume::PersistedAgentSession {
             source,
             agent: agent_label,
@@ -907,7 +958,13 @@ impl TerminalState {
         });
         let current_session = self.current_session_identity_for_persistence();
         Some(TerminalStateMutation {
-            effective_state_change: None,
+            effective_state_change: self.recompute_effective_state(
+                previous_agent_label,
+                previous_known_agent,
+                previous_state,
+                previous_presentation,
+                now,
+            ),
             session_ref_changed: previous_session != current_session,
         })
     }
@@ -1334,6 +1391,72 @@ mod tests {
     }
 
     #[test]
+    fn omp_resume_session_report_reanchors_full_lifecycle_authority() {
+        let mut terminal = test_terminal();
+        let old_session = test_session_path("omp-old.jsonl");
+        let new_session = test_session_path("omp-new.jsonl");
+        terminal.set_detected_state(Some(Agent::Omp), AgentState::Idle);
+        terminal.set_hook_authority_with_session_ref(
+            "herdr:omp".into(),
+            "omp".into(),
+            AgentState::Working,
+            None,
+            None,
+            crate::agent_resume::AgentSessionRef::path(old_session.clone()),
+            Some(10),
+        );
+
+        let session_report = terminal.set_agent_session_ref_for_session_start(
+            "herdr:omp".into(),
+            "omp".into(),
+            crate::agent_resume::AgentSessionRef::path(new_session.clone()),
+            Some(11),
+            Some("resume".into()),
+        );
+
+        assert!(session_report.is_some());
+        assert!(terminal.hook_authority.is_none());
+        assert_eq!(
+            terminal
+                .persisted_agent_session
+                .as_ref()
+                .unwrap()
+                .session_ref,
+            crate::agent_resume::AgentSessionRef::path(new_session.clone()).unwrap()
+        );
+
+        let blocked = terminal.set_hook_authority_with_session_ref(
+            "herdr:omp".into(),
+            "omp".into(),
+            AgentState::Blocked,
+            Some("waiting".into()),
+            None,
+            crate::agent_resume::AgentSessionRef::path(new_session.clone()),
+            Some(12),
+        );
+
+        assert!(blocked.is_some());
+        assert_eq!(terminal.state, AgentState::Blocked);
+        assert_eq!(
+            terminal.hook_authority.as_ref().unwrap().session_ref,
+            crate::agent_resume::AgentSessionRef::path(new_session)
+        );
+
+        let stale = terminal.set_hook_authority_with_session_ref(
+            "herdr:omp".into(),
+            "omp".into(),
+            AgentState::Working,
+            None,
+            None,
+            crate::agent_resume::AgentSessionRef::path(old_session),
+            Some(13),
+        );
+
+        assert!(stale.is_none());
+        assert_eq!(terminal.state, AgentState::Blocked);
+    }
+
+    #[test]
     fn process_exit_clears_matching_full_lifecycle_hook_authority() {
         let now = Instant::now();
         let mut terminal = test_terminal();