fix: restore omp sessions after same-pane restart

refs #614

4 files changed, 137 insertions, 11 deletions

diff --git a/src/agent_resume.rs b/src/agent_resume.rs
index 4a97f207..5012e991 100644
--- a/src/agent_resume.rs
+++ b/src/agent_resume.rs
@@ -60,7 +60,7 @@ pub fn session_ref_from_report(
         return None;
     }
 
-    if agent == "pi" {
+    if agent == "pi" || agent == "omp" {
         return _agent_session_path
             .and_then(AgentSessionRef::path)
             .or_else(|| agent_session_id.and_then(AgentSessionRef::id));
@@ -99,7 +99,7 @@ pub fn session_ref_from_snapshot(
         return None;
     }
     let session_ref = match (agent, kind) {
-        ("pi", AgentSessionRefKind::Path) => AgentSessionRef::path(value)?,
+        ("pi" | "omp", AgentSessionRefKind::Path) => AgentSessionRef::path(value)?,
         (_, AgentSessionRefKind::Id) => AgentSessionRef::id(value)?,
         _ => return None,
     };
@@ -141,6 +141,11 @@ pub fn plan(source: &str, agent: &str, session_ref: &AgentSessionRef) -> Option<
         ("herdr:pi", "pi", AgentSessionRefKind::Path | AgentSessionRefKind::Id) => {
             vec!["pi".into(), "--session".into(), session_ref.value.clone()]
         }
+        ("herdr:omp", "omp", AgentSessionRefKind::Path | AgentSessionRefKind::Id) => {
+            // omp resume is `-r, --resume=<value>` (ID prefix or path); it has no
+            // `--session` flag, unlike pi.
+            vec!["omp".into(), format!("--resume={}", session_ref.value)]
+        }
         ("herdr:hermes", "hermes", AgentSessionRefKind::Id) => {
             vec![
                 "hermes".into(),
@@ -198,6 +203,7 @@ fn is_official_agent_source(source: &str, agent: &str) -> bool {
             | ("herdr:devin", "devin")
             | ("herdr:droid", "droid")
             | ("herdr:kimi", "kimi")
+            | ("herdr:omp", "omp")
             | ("herdr:pi", "pi")
             | ("herdr:hermes", "hermes")
             | ("herdr:opencode", "opencode")
@@ -245,6 +251,7 @@ mod tests {
     #[test]
     fn planner_allows_supported_agents() {
         let pi_session = absolute_test_path("pi-session.jsonl");
+        let omp_session = absolute_test_path("omp-session.jsonl");
         assert_eq!(
             plan(
                 "herdr:claude",
@@ -317,6 +324,16 @@ mod tests {
         );
         assert_eq!(
             plan(
+                "herdr:omp",
+                "omp",
+                &AgentSessionRef::path(&omp_session).unwrap()
+            )
+            .unwrap()
+            .argv,
+            vec!["omp", format!("--resume={omp_session}").as_str()]
+        );
+        assert_eq!(
+            plan(
                 "herdr:hermes",
                 "hermes",
                 &AgentSessionRef::id("hermes-session").unwrap()
@@ -385,8 +402,9 @@ mod tests {
     }
 
     #[test]
-    fn report_ref_prefers_pi_path_and_validates_values() {
+    fn report_ref_prefers_pi_and_omp_paths_and_validates_values() {
         let pi_session = absolute_test_path("pi-session.jsonl");
+        let omp_session = absolute_test_path("omp-session.jsonl");
         let claude_session = absolute_test_path("claude-session");
         let copilot_session = absolute_test_path("copilot-session");
         let session_ref = session_ref_from_report(
@@ -405,6 +423,35 @@ mod tests {
                 .is_none()
         );
         assert!(session_ref_from_report("custom:pi", "pi", Some("pi-id".into()), None).is_none());
+
+        let session_ref = session_ref_from_report(
+            "herdr:omp",
+            "omp",
+            Some("omp-id".into()),
+            Some(omp_session.clone()),
+        )
+        .unwrap();
+        assert_eq!(session_ref.kind, AgentSessionRefKind::Path);
+        assert_eq!(session_ref.value, omp_session);
+
+        let session_ref =
+            session_ref_from_report("herdr:omp", "omp", Some("omp-id".into()), None).unwrap();
+        assert_eq!(session_ref.kind, AgentSessionRefKind::Id);
+        assert_eq!(session_ref.value, "omp-id");
+        let session_ref = session_ref_from_report(
+            "herdr:omp",
+            "omp",
+            Some("omp-id".into()),
+            Some("relative.jsonl".into()),
+        )
+        .unwrap();
+        assert_eq!(session_ref.kind, AgentSessionRefKind::Id);
+        assert_eq!(session_ref.value, "omp-id");
+        assert!(
+            session_ref_from_report("herdr:omp", "omp", None, Some("relative.jsonl".into()))
+                .is_none()
+        );
+
         assert!(
             session_ref_from_report("herdr:claude", "claude", None, Some(claude_session)).is_none()
         );
diff --git a/src/integration/assets/omp/herdr-agent-state.ts b/src/integration/assets/omp/herdr-agent-state.ts
index 9229436b..c2bec682 100644
--- a/src/integration/assets/omp/herdr-agent-state.ts
+++ b/src/integration/assets/omp/herdr-agent-state.ts
@@ -2,7 +2,7 @@
 // managed by herdr; reinstalling or updating the integration overwrites this file.
 // add custom hooks/plugins beside this file instead of editing it.
 // HERDR_INTEGRATION_ID=omp
-// HERDR_INTEGRATION_VERSION=2
+// HERDR_INTEGRATION_VERSION=3
 // @ts-nocheck
 
 import { createConnection } from "node:net";
@@ -53,12 +53,41 @@ const retryGraceMs = parseDurationEnv("HERDR_OMP_RETRY_GRACE_MS", 2500);
 const retryableErrorPattern =
   /overloaded|provider.?returned.?error|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server.?error|internal.?error|network.?error|connection.?error|connection.?refused|connection.?lost|websocket.?closed|websocket.?error|other side closed|fetch failed|upstream.?connect|reset before headers|socket hang up|ended without|http2 request did not get a response|timed? out|timeout|terminated|retry delay/i;
 let reportSeq = Date.now() * 1000;
+let currentAgentSessionId: string | undefined;
+let currentAgentSessionPath: string | undefined;
 
 function nextReportSeq(): number {
   reportSeq += 1;
   return reportSeq;
 }
 
+function updateSessionRef(ctx: any): void {
+  try {
+    const file = ctx?.sessionManager?.getSessionFile?.();
+    currentAgentSessionPath =
+      typeof file === "string" && file.startsWith("/") ? file : undefined;
+  } catch {
+    currentAgentSessionPath = undefined;
+  }
+
+  try {
+    const id = ctx?.sessionManager?.getSessionId?.();
+    currentAgentSessionId = typeof id === "string" && id.length > 0 ? id : undefined;
+  } catch {
+    currentAgentSessionId = undefined;
+  }
+}
+
+function withSessionRef(params: Record<string, unknown>): Record<string, unknown> {
+  if (currentAgentSessionPath) {
+    return { ...params, agent_session_path: currentAgentSessionPath };
+  }
+  if (currentAgentSessionId) {
+    return { ...params, agent_session_id: currentAgentSessionId };
+  }
+  return params;
+}
+
 function parseDurationEnv(name: string, fallback: number): number {
   const raw = process.env[name];
   if (!raw) {
@@ -75,14 +104,14 @@ function sendState(state: AgentState, message?: string, seq = nextReportSeq()):
   return sendRequest({
     id: `${source}:${Date.now()}:${Math.random().toString(36).slice(2)}`,
     method: "pane.report_agent",
-    params: {
+    params: withSessionRef({
       pane_id: paneId,
       source,
       agent: "omp",
       state,
       message,
       seq,
-    },
+    }),
   });
 }
 
@@ -201,9 +230,9 @@ export default function (pi) {
     return { state: "idle" as const, message: undefined };
   }
 
-  function publishState() {
+  function publishState(force = false) {
     const next = desiredState();
-    if (next.state === lastState && next.message === lastMessage) {
+    if (!force && next.state === lastState && next.message === lastMessage) {
       return;
     }
     lastState = next.state;
@@ -253,8 +282,9 @@ export default function (pi) {
     publishState();
   });
 
-  pi.on("session_start", () => {
-    publishState();
+  pi.on("session_start", (_event, ctx) => {
+    updateSessionRef(ctx);
+    publishState(true);
   });
 
   pi.on("agent_start", () => {
diff --git a/src/integration/mod.rs b/src/integration/mod.rs
index 234d1b10..d4c7265b 100644
--- a/src/integration/mod.rs
+++ b/src/integration/mod.rs
@@ -15,7 +15,7 @@ const PI_EXTENSION_ASSET: &str = include_str!("assets/pi/herdr-agent-state.ts");
 const PI_INTEGRATION_VERSION: u32 = 2;
 const OMP_EXTENSION_INSTALL_NAME: &str = "herdr-omp-agent-state.ts";
 const OMP_EXTENSION_ASSET: &str = include_str!("assets/omp/herdr-agent-state.ts");
-const OMP_INTEGRATION_VERSION: u32 = 2;
+const OMP_INTEGRATION_VERSION: u32 = 3;
 const PI_CODING_AGENT_DIR_ENV_VAR: &str = "PI_CODING_AGENT_DIR";
 const CLAUDE_HOOK_INSTALL_NAME: &str = if cfg!(windows) {
     "herdr-agent-state.ps1"
@@ -5969,6 +5969,9 @@ mod tests {
         assert!(PI_EXTENSION_ASSET.contains("agent_session_path: currentAgentSessionPath"));
         assert!(PI_EXTENSION_ASSET.contains("agent_session_id: currentAgentSessionId"));
         assert!(PI_EXTENSION_ASSET.contains("publishState(true)"));
+        assert!(OMP_EXTENSION_ASSET.contains("agent_session_path: currentAgentSessionPath"));
+        assert!(OMP_EXTENSION_ASSET.contains("agent_session_id: currentAgentSessionId"));
+        assert!(OMP_EXTENSION_ASSET.contains("publishState(true)"));
         assert!(CLAUDE_HOOK_ASSET.contains("agent_session_id"));
         assert!(CLAUDE_HOOK_ASSET.contains("agent_session_path"));
         assert!(CLAUDE_HOOK_ASSET.contains("session_start_source"));
diff --git a/src/terminal/state.rs b/src/terminal/state.rs
index 9b9b6d43..07e84354 100644
--- a/src/terminal/state.rs
+++ b/src/terminal/state.rs
@@ -1493,6 +1493,52 @@ mod tests {
         assert_eq!(terminal.state, AgentState::Working);
     }
 
+    // omp has no Agent identity (the process is pi), so its release suppression
+    // can only be cleared by a fresh session ref, never by process detection.
+    // Regression for #614: a same-pane restart must reacquire lifecycle authority.
+    #[test]
+    fn omp_reacquires_full_lifecycle_hook_after_release_with_fresh_session_ref() {
+        let mut terminal = test_terminal();
+        terminal.set_hook_authority_with_session_ref(
+            "herdr:omp".into(),
+            "omp".into(),
+            AgentState::Working,
+            None,
+            None,
+            crate::agent_resume::AgentSessionRef::id("omp-old"),
+            Some(20),
+        );
+        terminal.release_agent("herdr:omp", "omp", Some(21));
+
+        // A late report from the released run keeps its old session ref and stays
+        // suppressed, so a just-exited omp cannot resurrect the pane.
+        let stale = terminal.set_hook_authority_with_session_ref(
+            "herdr:omp".into(),
+            "omp".into(),
+            AgentState::Working,
+            None,
+            None,
+            crate::agent_resume::AgentSessionRef::id("omp-old"),
+            Some(22),
+        );
+        assert!(stale.is_none());
+        assert!(terminal.hook_authority.is_none());
+
+        // A fresh omp run carries a new session ref and reacquires authority.
+        let fresh = terminal.set_hook_authority_with_session_ref(
+            "herdr:omp".into(),
+            "omp".into(),
+            AgentState::Working,
+            None,
+            None,
+            crate::agent_resume::AgentSessionRef::id("omp-new"),
+            Some(23),
+        );
+        assert!(fresh.is_some());
+        assert!(terminal.hook_authority.is_some());
+        assert_eq!(terminal.state, AgentState::Working);
+    }
+
     #[test]
     fn visible_blocker_overrides_non_blocked_hook_for_same_agent() {
         let mut terminal = test_terminal();