fix: avoid pi release on internal session reloads (#800)

* fix: avoid pi release on internal session reloads

* fix: align omp session shutdown release with pi

---------

Co-authored-by: Can Celik <ogulcancelik@gmail.com>

3 files changed, 120 insertions, 10 deletions

diff --git a/src/integration/assets/omp/herdr-agent-state.ts b/src/integration/assets/omp/herdr-agent-state.ts
index ba97b1c7..e0cfab5e 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=3
+// HERDR_INTEGRATION_VERSION=4
 // @ts-nocheck
 
 import { createConnection } from "node:net";
@@ -157,6 +157,16 @@ function releaseAgent(): Promise<void> {
   });
 }
 
+function shouldReleaseOnSessionShutdown(event: any): boolean {
+  // OMP tears down and rebinds extension runtimes for internal lifecycle actions
+  // such as /reload, /new, /resume, and /fork. Those do not mean the pane's
+  // agent process has exited, and releasing hook authority there can suppress
+  // legitimate reports from the replacement runtime. Only a user/process quit
+  // should release Herdr's full-lifecycle authority.
+  const reason = event?.reason;
+  return reason === "quit";
+}
+
 let sendInFlight = false;
 let queuedState: QueuedState | undefined;
 
@@ -325,10 +335,12 @@ export default function (pi) {
     publishState(true);
   });
 
-  pi.on("agent_start", () => {
+  pi.on("agent_start", (_event, ctx) => {
     if (!rootSession) {
       return;
     }
+    updateSessionRef(ctx);
+    void reportSession();
     clearPendingTimers();
     clearFailureState();
     agentActive = true;
@@ -357,11 +369,13 @@ export default function (pi) {
     scheduleIdle();
   });
 
-  pi.on("session_shutdown", async () => {
+  pi.on("session_shutdown", async (event) => {
     if (!rootSession) {
       return;
     }
     clearPendingTimers();
-    await releaseAgent();
+    if (shouldReleaseOnSessionShutdown(event)) {
+      await releaseAgent();
+    }
   });
 }
diff --git a/src/integration/assets/pi/herdr-agent-state.ts b/src/integration/assets/pi/herdr-agent-state.ts
index 6c9e1cae..4bee2763 100644
--- a/src/integration/assets/pi/herdr-agent-state.ts
+++ b/src/integration/assets/pi/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=pi
-// HERDR_INTEGRATION_VERSION=3
+// HERDR_INTEGRATION_VERSION=4
 // @ts-nocheck
 
 import { createConnection } from "node:net";
@@ -157,6 +157,16 @@ function releaseAgent(): Promise<void> {
   });
 }
 
+function shouldReleaseOnSessionShutdown(event: any): boolean {
+  // Pi tears down and rebinds extension runtimes for internal lifecycle actions
+  // such as /reload, /new, /resume, and /fork. Those do not mean the pane's
+  // agent process has exited, and releasing hook authority there can suppress
+  // legitimate reports from the replacement runtime. Only a user/process quit
+  // should release Herdr's full-lifecycle authority.
+  const reason = event?.reason;
+  return reason === "quit";
+}
+
 let sendInFlight = false;
 let queuedState: QueuedState | undefined;
 
@@ -325,10 +335,12 @@ export default function (pi) {
     publishState(true);
   });
 
-  pi.on("agent_start", () => {
+  pi.on("agent_start", (_event, ctx) => {
     if (!rootSession) {
       return;
     }
+    updateSessionRef(ctx);
+    void reportSession();
     clearPendingTimers();
     clearFailureState();
     agentActive = true;
@@ -357,11 +369,13 @@ export default function (pi) {
     scheduleIdle();
   });
 
-  pi.on("session_shutdown", async () => {
+  pi.on("session_shutdown", async (event) => {
     if (!rootSession) {
       return;
     }
     clearPendingTimers();
-    await releaseAgent();
+    if (shouldReleaseOnSessionShutdown(event)) {
+      await releaseAgent();
+    }
   });
 }
diff --git a/src/integration/mod.rs b/src/integration/mod.rs
index 06561e3d..287ab89b 100644
--- a/src/integration/mod.rs
+++ b/src/integration/mod.rs
@@ -12,10 +12,10 @@ pub(crate) const HERDR_TAB_ID_ENV_VAR: &str = "HERDR_TAB_ID";
 pub(crate) const HERDR_WORKSPACE_ID_ENV_VAR: &str = "HERDR_WORKSPACE_ID";
 const PI_EXTENSION_INSTALL_NAME: &str = "herdr-agent-state.ts";
 const PI_EXTENSION_ASSET: &str = include_str!("assets/pi/herdr-agent-state.ts");
-const PI_INTEGRATION_VERSION: u32 = 3;
+const PI_INTEGRATION_VERSION: u32 = 4;
 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 = 3;
+const OMP_INTEGRATION_VERSION: u32 = 4;
 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"
@@ -6074,6 +6074,88 @@ mod tests {
     }
 
     #[test]
+    fn pi_extension_releases_only_for_quit_session_shutdown() {
+        let release_policy = PI_EXTENSION_ASSET
+            .find("function shouldReleaseOnSessionShutdown")
+            .expect("pi extension should centralize session shutdown release policy");
+        let quit_check = PI_EXTENSION_ASSET
+            .find("reason === \"quit\"")
+            .expect("pi extension should release only for true quit shutdowns");
+        let shutdown_handler = PI_EXTENSION_ASSET
+            .find("pi.on(\"session_shutdown\", async (event)")
+            .expect("pi extension should inspect the session_shutdown event");
+        let guarded_release = PI_EXTENSION_ASSET[shutdown_handler..]
+            .find("if (shouldReleaseOnSessionShutdown(event))")
+            .expect("pi extension should guard releaseAgent by shutdown reason");
+
+        assert!(release_policy < shutdown_handler);
+        assert!(release_policy < quit_check);
+        assert!(quit_check < shutdown_handler);
+        assert!(guarded_release > 0);
+    }
+
+    #[test]
+    fn pi_extension_refreshes_session_ref_before_agent_start_state() {
+        let agent_start = PI_EXTENSION_ASSET
+            .find("pi.on(\"agent_start\", (_event, ctx)")
+            .expect("pi extension should receive agent_start context");
+        let handler = &PI_EXTENSION_ASSET[agent_start..];
+        let update_session = handler
+            .find("updateSessionRef(ctx);")
+            .expect("pi extension should refresh the active session on agent_start");
+        let report_session = handler
+            .find("void reportSession();")
+            .expect("pi extension should report the refreshed session before state");
+        let publish_state = handler
+            .find("publishState();")
+            .expect("pi extension should publish working state after refreshing session");
+
+        assert!(update_session < report_session);
+        assert!(report_session < publish_state);
+    }
+
+    #[test]
+    fn omp_extension_releases_only_for_quit_session_shutdown() {
+        let release_policy = OMP_EXTENSION_ASSET
+            .find("function shouldReleaseOnSessionShutdown")
+            .expect("omp extension should centralize session shutdown release policy");
+        let quit_check = OMP_EXTENSION_ASSET
+            .find("reason === \"quit\"")
+            .expect("omp extension should release only for true quit shutdowns");
+        let shutdown_handler = OMP_EXTENSION_ASSET
+            .find("pi.on(\"session_shutdown\", async (event)")
+            .expect("omp extension should inspect the session_shutdown event");
+        let guarded_release = OMP_EXTENSION_ASSET[shutdown_handler..]
+            .find("if (shouldReleaseOnSessionShutdown(event))")
+            .expect("omp extension should guard releaseAgent by shutdown reason");
+
+        assert!(release_policy < shutdown_handler);
+        assert!(release_policy < quit_check);
+        assert!(quit_check < shutdown_handler);
+        assert!(guarded_release > 0);
+    }
+
+    #[test]
+    fn omp_extension_refreshes_session_ref_before_agent_start_state() {
+        let agent_start = OMP_EXTENSION_ASSET
+            .find("pi.on(\"agent_start\", (_event, ctx)")
+            .expect("omp extension should receive agent_start context");
+        let handler = &OMP_EXTENSION_ASSET[agent_start..];
+        let update_session = handler
+            .find("updateSessionRef(ctx);")
+            .expect("omp extension should refresh the active session on agent_start");
+        let report_session = handler
+            .find("void reportSession();")
+            .expect("omp extension should report the refreshed session before state");
+        let publish_state = handler
+            .find("publishState();")
+            .expect("omp extension should publish working state after refreshing session");
+
+        assert!(update_session < report_session);
+        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\"")