fix: reanchor pi status after session replacement (#1189)

* fix: reanchor pi status after session replacement

refs #943

* fix: guard active pi sessions from nested startup

refs #943

* fix: serialize pi session replacement reports

refs #943

src/integration/assets/herdr-agent-state.test.ts

@@ -48,42 +48,13 @@ function importFresh(modulePath: string) {
   return import(`${modulePath}?test=${importCounter}`);
 }
 
-for (const integration of integrations) {
-  test(`${integration.name} reload preserves working state when the agent is active`, async () => {
-    const recordingSocketPath = join(
-      tmpdir(),
-      `herdr-${integration.name.toLowerCase().replaceAll(" ", "-")}-${process.pid}.sock`,
-    );
-    socketPath = recordingSocketPath;
-    await rm(recordingSocketPath, { force: true });
+type Handler = (event: unknown, context: unknown) => unknown;
 
-    const requests: unknown[] = [];
-    const recordingServer = createServer((socket) => {
-      let input = "";
-      socket.setEncoding("utf8");
-      socket.on("data", (chunk) => {
-        input += chunk;
-        const newline = input.indexOf("\n");
-        if (newline === -1) {
-          return;
-        }
-        requests.push(JSON.parse(input.slice(0, newline)));
-        socket.end("{}\n");
-      });
-    });
-    server = recordingServer;
-    await new Promise<void>((resolve, reject) => {
-      recordingServer.once("error", reject);
-      recordingServer.listen(recordingSocketPath, resolve);
-    });
-
-    process.env.HERDR_ENV = "1";
-    process.env.HERDR_SOCKET_PATH = recordingSocketPath;
-    process.env.HERDR_PANE_ID = "test:p1";
-
-    type Handler = (event: unknown, context: unknown) => unknown;
-    const handlers = new Map<string, Handler>();
-    const pi = {
+function createExtensionHarness() {
+  const handlers = new Map<string, Handler>();
+  return {
+    handlers,
+    pi: {
       on(event: string, handler: Handler) {
         handlers.set(event, handler);
       },
@@ -92,7 +63,50 @@ for (const integration of integrations) {
           return () => {};
         },
       },
-    };
+    },
+  };
+}
+
+function configureIntegrationEnvironment(recordingSocketPath: string) {
+  process.env.HERDR_ENV = "1";
+  process.env.HERDR_SOCKET_PATH = recordingSocketPath;
+  process.env.HERDR_PANE_ID = "test:p1";
+}
+
+async function startRecordingServer(name: string): Promise<unknown[]> {
+  const recordingSocketPath = join(tmpdir(), `herdr-${name}-${process.pid}.sock`);
+  socketPath = recordingSocketPath;
+  await rm(recordingSocketPath, { force: true });
+
+  const requests: unknown[] = [];
+  const recordingServer = createServer((socket) => {
+    let input = "";
+    socket.setEncoding("utf8");
+    socket.on("data", (chunk) => {
+      input += chunk;
+      const newline = input.indexOf("\n");
+      if (newline === -1) {
+        return;
+      }
+      requests.push(JSON.parse(input.slice(0, newline)));
+      socket.end("{}\n");
+    });
+  });
+  server = recordingServer;
+  await new Promise<void>((resolve, reject) => {
+    recordingServer.once("error", reject);
+    recordingServer.listen(recordingSocketPath, resolve);
+  });
+  configureIntegrationEnvironment(recordingSocketPath);
+  return requests;
+}
+
+for (const integration of integrations) {
+  test(`${integration.name} reload preserves working state when the agent is active`, async () => {
+    const requests = await startRecordingServer(
+      integration.name.toLowerCase().replaceAll(" ", "-"),
+    );
+    const { handlers, pi } = createExtensionHarness();
 
     const { default: install } = await importFresh(integration.modulePath);
     install(pi);
@@ -133,6 +147,115 @@ for (const integration of integrations) {
   });
 }
 
+test("Pi reports the session replacement source", async () => {
+  const requests = await startRecordingServer("pi-session-source");
+  const { handlers, pi } = createExtensionHarness();
+
+  const { default: install } = await importFresh("./pi/herdr-agent-state.ts");
+  install(pi);
+
+  const sessionStart = handlers.get("session_start");
+  expect(sessionStart).toBeDefined();
+  await sessionStart?.(
+    { reason: "new" },
+    {
+      hasUI: true,
+      isIdle: () => true,
+      sessionManager: {
+        getSessionFile: () => "/tmp/pi-new.jsonl",
+        getSessionId: () => "pi-new",
+      },
+    },
+  );
+
+  const reportedSession = () =>
+    requests.find((request) => isRecord(request) && request.method === "pane.report_agent_session");
+  const deadline = Date.now() + 1_000;
+  while (Date.now() < deadline && reportedSession() === undefined) {
+    await Bun.sleep(5);
+  }
+
+  const request = reportedSession();
+  expect(request).toBeDefined();
+  expect(isRecord(request) && isRecord(request.params) ? request.params.session_start_source : null)
+    .toBe("new");
+});
+
+test("Pi waits for a replacement session report before publishing state", async () => {
+  const recordingSocketPath = join(tmpdir(), `herdr-pi-session-order-${process.pid}.sock`);
+  socketPath = recordingSocketPath;
+  await rm(recordingSocketPath, { force: true });
+
+  const requests: unknown[] = [];
+  let acknowledgeSessionReport: (() => void) | undefined;
+  const recordingServer = createServer((socket) => {
+    let input = "";
+    socket.setEncoding("utf8");
+    socket.on("data", (chunk) => {
+      input += chunk;
+      const newline = input.indexOf("\n");
+      if (newline === -1) {
+        return;
+      }
+      const request = JSON.parse(input.slice(0, newline));
+      requests.push(request);
+      if (isRecord(request) && request.method === "pane.report_agent_session") {
+        acknowledgeSessionReport = () => socket.end("{}\n");
+        return;
+      }
+      socket.end("{}\n");
+    });
+  });
+  server = recordingServer;
+  await new Promise<void>((resolve, reject) => {
+    recordingServer.once("error", reject);
+    recordingServer.listen(recordingSocketPath, resolve);
+  });
+
+  configureIntegrationEnvironment(recordingSocketPath);
+  const { handlers, pi } = createExtensionHarness();
+  const { default: install } = await importFresh("./pi/herdr-agent-state.ts");
+  install(pi);
+
+  const sessionStart = handlers.get("session_start");
+  expect(sessionStart).toBeDefined();
+  const sessionStartResult = sessionStart?.(
+    { reason: "new" },
+    {
+      hasUI: true,
+      isIdle: () => false,
+      sessionManager: {
+        getSessionFile: () => "/tmp/pi-new.jsonl",
+        getSessionId: () => "pi-new",
+      },
+    },
+  );
+
+  const deadline = Date.now() + 1_000;
+  while (Date.now() < deadline && acknowledgeSessionReport === undefined) {
+    await Bun.sleep(5);
+  }
+  expect(acknowledgeSessionReport).toBeDefined();
+  expect(
+    requests.some((request) => isRecord(request) && request.method === "pane.report_agent"),
+  ).toBe(false);
+
+  acknowledgeSessionReport?.();
+  await sessionStartResult;
+
+  const stateDeadline = Date.now() + 1_000;
+  while (
+    Date.now() < stateDeadline &&
+    !requests.some((request) => isRecord(request) && request.method === "pane.report_agent")
+  ) {
+    await Bun.sleep(5);
+  }
+  expect(requests.map((request) => (isRecord(request) ? request.method : undefined))).toEqual([
+    "pane.report_agent_session",
+    "pane.report_agent",
+  ]);
+});
+
 test("Pi retries working state after an unanswered socket attempt", async () => {
   const recordingSocketPath = join(tmpdir(), `herdr-pi-retry-${process.pid}.sock`);
   socketPath = recordingSocketPath;
@@ -167,22 +290,8 @@ test("Pi retries working state after an unanswered socket attempt", async () =>
     recordingServer.listen(recordingSocketPath, resolve);
   });
 
-  process.env.HERDR_ENV = "1";
-  process.env.HERDR_SOCKET_PATH = recordingSocketPath;
-  process.env.HERDR_PANE_ID = "test:p1";
-
-  type Handler = (event: unknown, context: unknown) => unknown;
-  const handlers = new Map<string, Handler>();
-  const pi = {
-    on(event: string, handler: Handler) {
-      handlers.set(event, handler);
-    },
-    events: {
-      on() {
-        return () => {};
-      },
-    },
-  };
+  configureIntegrationEnvironment(recordingSocketPath);
+  const { handlers, pi } = createExtensionHarness();
 
   const { default: install } = await importFresh("./pi/herdr-agent-state.ts");
   install(pi);

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=4
+// HERDR_INTEGRATION_VERSION=5
 // @ts-nocheck
 
 import { createConnection } from "node:net";
@@ -121,7 +121,7 @@ function currentSessionRef(): Record<string, unknown> | undefined {
   return undefined;
 }
 
-function reportSession(): Promise<void> {
+function reportSession(sessionStartSource?: string): Promise<void> {
   const sessionRef = currentSessionRef();
   if (!sessionRef) {
     return Promise.resolve();
@@ -135,6 +135,7 @@ function reportSession(): Promise<void> {
       source,
       agent: "pi",
       seq: nextReportSeq(),
+      session_start_source: sessionStartSource,
       ...sessionRef,
     },
   });
@@ -336,13 +337,13 @@ export default function (pi) {
     publishState();
   });
 
-  pi.on("session_start", (_event, ctx) => {
+  pi.on("session_start", async (event, ctx) => {
     if (ctx?.hasUI !== true) {
       return;
     }
     rootSession = true;
     updateSessionRef(ctx);
-    void reportSession();
+    await reportSession(event?.reason);
     // A reload can replace this extension mid-run without emitting another agent_start.
     agentActive = ctx?.isIdle?.() === false;
     publishState(true);

src/integration/mod.rs

@@ -22,7 +22,7 @@ pub(crate) use types::{IntegrationRecommendation, IntegrationStatus, Integration
 
 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 = 4;
+const PI_INTEGRATION_VERSION: u32 = 5;
 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 = 4;

src/integration/tests.rs

@@ -732,6 +732,36 @@ fn outdated_integrations_treat_missing_version_marker_as_legacy() {
     let _ = fs::remove_dir_all(base);
 }
 
+#[test]
+fn outdated_integrations_detect_previous_pi_version() {
+    let _lock = integration_env_lock();
+    let base = unique_base();
+    let home = base.join("home");
+    let ext_dir = home.join(".pi/agent/extensions");
+    fs::create_dir_all(&ext_dir).unwrap();
+    let extension_path = ext_dir.join(PI_EXTENSION_INSTALL_NAME);
+    fs::write(
+        &extension_path,
+        "// HERDR_INTEGRATION_ID=pi\n// HERDR_INTEGRATION_VERSION=4\n",
+    )
+    .unwrap();
+    std::env::set_var("HOME", &home);
+
+    let outdated = outdated_installed_integrations();
+
+    assert_eq!(outdated.len(), 1);
+    assert_eq!(
+        outdated[0].target,
+        crate::api::schema::IntegrationTarget::Pi
+    );
+    assert_eq!(outdated[0].path, extension_path);
+    assert_eq!(outdated[0].installed_version, Some(4));
+    assert_eq!(outdated[0].expected_version, PI_INTEGRATION_VERSION);
+
+    std::env::remove_var("HOME");
+    let _ = fs::remove_dir_all(base);
+}
+
 #[test]
 fn outdated_integrations_accept_current_version_marker() {
     let _lock = integration_env_lock();
@@ -2482,6 +2512,48 @@ fn install_hermes_errors_when_config_dir_missing() {
     let _ = fs::remove_dir_all(base);
 }
 
+#[test]
+fn bundled_integration_asset_versions_match_expected_versions() {
+    for (name, asset, expected_version) in [
+        ("pi", PI_EXTENSION_ASSET, PI_INTEGRATION_VERSION),
+        ("omp", OMP_EXTENSION_ASSET, OMP_INTEGRATION_VERSION),
+        ("claude", CLAUDE_HOOK_ASSET, CLAUDE_INTEGRATION_VERSION),
+        ("codex", CODEX_HOOK_ASSET, CODEX_INTEGRATION_VERSION),
+        ("kimi", KIMI_HOOK_ASSET, KIMI_INTEGRATION_VERSION),
+        ("copilot", COPILOT_HOOK_ASSET, COPILOT_INTEGRATION_VERSION),
+        ("devin", DEVIN_HOOK_ASSET, DEVIN_INTEGRATION_VERSION),
+        ("droid", DROID_HOOK_ASSET, DROID_INTEGRATION_VERSION),
+        (
+            "opencode",
+            OPENCODE_PLUGIN_ASSET,
+            OPENCODE_INTEGRATION_VERSION,
+        ),
+        ("kilo", KILO_PLUGIN_ASSET, KILO_INTEGRATION_VERSION),
+        (
+            "hermes",
+            HERMES_PLUGIN_INIT_ASSET,
+            HERMES_INTEGRATION_VERSION,
+        ),
+        (
+            "qodercli",
+            QODERCLI_HOOK_ASSET,
+            QODERCLI_INTEGRATION_VERSION,
+        ),
+        ("cursor", CURSOR_HOOK_ASSET, CURSOR_INTEGRATION_VERSION),
+        (
+            "mastracode",
+            MASTRACODE_HOOK_ASSET,
+            MASTRACODE_INTEGRATION_VERSION,
+        ),
+    ] {
+        assert_eq!(
+            parse_integration_version(asset),
+            Some(expected_version),
+            "{name} asset version must match its integration version constant"
+        );
+    }
+}
+
 #[test]
 fn bundled_integration_assets_report_session_refs() {
     assert!(PI_EXTENSION_ASSET.contains("agent_session_path"));

src/terminal/state.rs

@@ -797,6 +797,26 @@ impl TerminalState {
         }
     }
 
+    fn forget_stale_full_lifecycle_hook_session(
+        &mut self,
+        source: &str,
+        agent_label: &str,
+        session_ref: &crate::agent_resume::AgentSessionRef,
+    ) {
+        let remove_source = self
+            .stale_full_lifecycle_hook_sessions
+            .get_mut(source)
+            .is_some_and(|stale_sessions| {
+                stale_sessions.retain(|stale| {
+                    stale.agent_label != agent_label || &stale.session_ref != session_ref
+                });
+                stale_sessions.is_empty()
+            });
+        if remove_source {
+            self.stale_full_lifecycle_hook_sessions.remove(source);
+        }
+    }
+
     fn detected_state_observed_before_release_suppression(
         &self,
         detected_agent: Option<Agent>,
@@ -912,6 +932,7 @@ impl TerminalState {
                 "codex",
                 Some("startup" | "clear" | "resume" | "compact")
             ) | ("herdr:opencode", "opencode", Some("new"))
+                | ("herdr:pi", "pi", Some("new" | "resume" | "fork"))
                 | (
                     "herdr:omp",
                     "omp",
@@ -1001,6 +1022,9 @@ 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 session_replacement_allowed || foreground_takeover_allowed {
+            self.forget_stale_full_lifecycle_hook_session(&source, &agent_label, &session_ref);
+        }
         if let Some(replaced_hook_session) = replaced_hook_session {
             self.remember_stale_full_lifecycle_hook_session(
                 source.clone(),
@@ -1488,6 +1512,208 @@ mod tests {
         }
     }
 
+    #[test]
+    fn pi_session_replacement_reports_reanchor_full_lifecycle_authority() {
+        for reason in ["new", "resume", "fork"] {
+            let mut terminal = test_terminal();
+            let old_session = test_session_path(&format!("pi-{reason}-old.jsonl"));
+            let new_session = test_session_path(&format!("pi-{reason}-new.jsonl"));
+            terminal.set_detected_state(Some(Agent::Pi), AgentState::Idle);
+            terminal.set_hook_authority_with_session_ref(
+                "herdr:pi".into(),
+                "pi".into(),
+                AgentState::Idle,
+                None,
+                None,
+                crate::agent_resume::AgentSessionRef::path(old_session),
+                Some(10),
+            );
+
+            let session_report = terminal.set_agent_session_ref_for_session_start(
+                "herdr:pi".into(),
+                "pi".into(),
+                crate::agent_resume::AgentSessionRef::path(new_session.clone()),
+                Some(11),
+                Some(reason.into()),
+            );
+
+            assert!(
+                session_report.is_some(),
+                "{reason} should replace the previous Pi session"
+            );
+            assert!(terminal.hook_authority.is_none());
+
+            let working = terminal.set_hook_authority_with_session_ref(
+                "herdr:pi".into(),
+                "pi".into(),
+                AgentState::Working,
+                None,
+                None,
+                crate::agent_resume::AgentSessionRef::path(new_session.clone()),
+                Some(12),
+            );
+
+            assert!(
+                working.is_some(),
+                "{reason} should accept working for the replacement session"
+            );
+            assert_eq!(terminal.state, AgentState::Working);
+            assert_eq!(
+                terminal.hook_authority.as_ref().unwrap().session_ref,
+                crate::agent_resume::AgentSessionRef::path(new_session)
+            );
+        }
+    }
+
+    #[test]
+    fn pi_resume_reactivates_a_previously_stale_session() {
+        let mut terminal = test_terminal();
+        let session_a = test_session_path("pi-session-a.jsonl");
+        let session_b = test_session_path("pi-session-b.jsonl");
+        terminal.set_detected_state(Some(Agent::Pi), AgentState::Idle);
+        terminal.set_hook_authority_with_session_ref(
+            "herdr:pi".into(),
+            "pi".into(),
+            AgentState::Idle,
+            None,
+            None,
+            crate::agent_resume::AgentSessionRef::path(session_a.clone()),
+            Some(10),
+        );
+
+        terminal.set_agent_session_ref_for_session_start(
+            "herdr:pi".into(),
+            "pi".into(),
+            crate::agent_resume::AgentSessionRef::path(session_b.clone()),
+            Some(11),
+            Some("new".into()),
+        );
+        terminal.set_hook_authority_with_session_ref(
+            "herdr:pi".into(),
+            "pi".into(),
+            AgentState::Idle,
+            None,
+            None,
+            crate::agent_resume::AgentSessionRef::path(session_b.clone()),
+            Some(12),
+        );
+
+        let resumed = terminal.set_agent_session_ref_for_session_start(
+            "herdr:pi".into(),
+            "pi".into(),
+            crate::agent_resume::AgentSessionRef::path(session_a.clone()),
+            Some(13),
+            Some("resume".into()),
+        );
+        let working = terminal.set_hook_authority_with_session_ref(
+            "herdr:pi".into(),
+            "pi".into(),
+            AgentState::Working,
+            None,
+            None,
+            crate::agent_resume::AgentSessionRef::path(session_a.clone()),
+            Some(14),
+        );
+
+        assert!(resumed.is_some());
+        assert!(working.is_some());
+        assert_eq!(terminal.state, AgentState::Working);
+        assert_eq!(
+            terminal.hook_authority.as_ref().unwrap().session_ref,
+            crate::agent_resume::AgentSessionRef::path(session_a)
+        );
+
+        let late_session_b = terminal.set_hook_authority_with_session_ref(
+            "herdr:pi".into(),
+            "pi".into(),
+            AgentState::Idle,
+            None,
+            None,
+            crate::agent_resume::AgentSessionRef::path(session_b),
+            Some(15),
+        );
+        assert!(late_session_b.is_none());
+        assert_eq!(terminal.state, AgentState::Working);
+    }
+
+    #[test]
+    fn pi_startup_adopts_persisted_session_without_live_authority() {
+        let mut terminal = test_terminal();
+        let old_session = test_session_path("pi-startup-old.jsonl");
+        let new_session = test_session_path("pi-startup-new.jsonl");
+        terminal.set_detected_state(Some(Agent::Pi), AgentState::Idle);
+        terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
+            source: "herdr:pi".into(),
+            agent: "pi".into(),
+            session_ref: crate::agent_resume::AgentSessionRef::path(old_session)
+                .expect("test session path should be valid"),
+        });
+
+        let startup = terminal.set_agent_session_ref_for_session_start(
+            "herdr:pi".into(),
+            "pi".into(),
+            crate::agent_resume::AgentSessionRef::path(new_session.clone()),
+            Some(11),
+            Some("startup".into()),
+        );
+
+        assert!(startup.is_some());
+        assert_eq!(
+            terminal.current_session_identity_for_persistence(),
+            Some((
+                "herdr:pi".into(),
+                "pi".into(),
+                crate::agent_resume::AgentSessionRefKind::Path,
+                new_session,
+            ))
+        );
+    }
+
+    #[test]
+    fn pi_non_replacement_reports_preserve_full_lifecycle_authority() {
+        for reason in [None, Some("reload"), Some("startup")] {
+            let mut terminal = test_terminal();
+            let old_session = test_session_path("pi-current.jsonl");
+            let new_session = test_session_path("pi-unexpected.jsonl");
+            terminal.set_detected_state(Some(Agent::Pi), AgentState::Idle);
+            terminal.set_hook_authority_with_session_ref(
+                "herdr:pi".into(),
+                "pi".into(),
+                AgentState::Idle,
+                None,
+                None,
+                crate::agent_resume::AgentSessionRef::path(old_session.clone()),
+                Some(10),
+            );
+
+            let session_report = terminal.set_agent_session_ref_for_session_start(
+                "herdr:pi".into(),
+                "pi".into(),
+                crate::agent_resume::AgentSessionRef::path(new_session.clone()),
+                Some(11),
+                reason.map(str::to_string),
+            );
+            let working = terminal.set_hook_authority_with_session_ref(
+                "herdr:pi".into(),
+                "pi".into(),
+                AgentState::Working,
+                None,
+                None,
+                crate::agent_resume::AgentSessionRef::path(new_session),
+                Some(12),
+            );
+
+            assert!(session_report.is_none());
+            assert!(working.is_none());
+            assert_eq!(terminal.state, AgentState::Idle);
+            assert_eq!(
+                terminal.hook_authority.as_ref().unwrap().session_ref,
+                crate::agent_resume::AgentSessionRef::path(old_session),
+                "{reason:?} must not replace the current Pi session"
+            );
+        }
+    }
+
     #[test]
     fn omp_resume_session_report_reanchors_full_lifecycle_authority() {
         let mut terminal = test_terminal();

tests/cli_wrapper.rs

@@ -1524,7 +1524,7 @@ fn integration_commands_run_locally_when_server_is_missing() {
         .unwrap();
     assert_eq!(integration_status.status.code(), Some(0));
     let status_stdout = String::from_utf8_lossy(&integration_status.stdout);
-    assert!(status_stdout.contains("pi: current (v4)"));
+    assert!(status_stdout.contains("pi: current (v5)"));
     assert!(status_stdout.contains("claude: not installed"));
 
     let integration_uninstall = Command::new(env!("CARGO_BIN_EXE_herdr"))