fix: protect root agent restore sessions

refs #712

1 files changed, 50 insertions, 24 deletions

diff --git a/src/integration/assets/pi/herdr-agent-state.ts b/src/integration/assets/pi/herdr-agent-state.ts
index 0e620f1f..9af12cfd 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=2
+// HERDR_INTEGRATION_VERSION=3
 // @ts-nocheck
 
 import { createConnection } from "node:net";
@@ -100,6 +100,35 @@ function withSessionRef(params: Record<string, unknown>): Record<string, unknown
   return params;
 }
 
+function currentSessionRef(): Record<string, unknown> | undefined {
+  if (currentAgentSessionPath) {
+    return { agent_session_path: currentAgentSessionPath };
+  }
+  if (currentAgentSessionId) {
+    return { agent_session_id: currentAgentSessionId };
+  }
+  return undefined;
+}
+
+function reportSession(): Promise<void> {
+  const sessionRef = currentSessionRef();
+  if (!sessionRef) {
+    return Promise.resolve();
+  }
+
+  return sendRequest({
+    id: `${source}:session:${Date.now()}:${Math.random().toString(36).slice(2)}`,
+    method: "pane.report_agent_session",
+    params: {
+      pane_id: paneId,
+      source,
+      agent: "pi",
+      seq: nextReportSeq(),
+      ...sessionRef,
+    },
+  });
+}
+
 function sendState(state: AgentState, message?: string, seq = nextReportSeq()): Promise<void> {
   return sendRequest({
     id: `${source}:${Date.now()}:${Math.random().toString(36).slice(2)}`,
@@ -169,19 +198,6 @@ function retryableErrorMessage(event: any): string | undefined {
   return errorMessage || "retryable provider error";
 }
 
-function releaseAgent(): Promise<void> {
-  return sendRequest({
-    id: `${source}:release:${Date.now()}:${Math.random().toString(36).slice(2)}`,
-    method: "pane.release_agent",
-    params: {
-      pane_id: paneId,
-      source,
-      agent: "pi",
-      seq: nextReportSeq(),
-    },
-  });
-}
-
 export default function (pi) {
   if (!enabled()) {
     return;
@@ -197,6 +213,7 @@ export default function (pi) {
   let lastMessage: string | undefined;
   let idleTimer: ReturnType<typeof setTimeout> | undefined;
   let retryTimer: ReturnType<typeof setTimeout> | undefined;
+  let rootSession = false;
 
   function clearTimer(timer: ReturnType<typeof setTimeout> | undefined) {
     if (timer) {
@@ -250,11 +267,6 @@ export default function (pi) {
     idleTimer.unref?.();
   }
 
-  pi.on("session_start", (_event, ctx) => {
-    updateSessionRef(ctx);
-    publishState(true);
-  });
-
   function holdForRetry(message: string) {
     clearPendingTimers();
     retryHoldActive = true;
@@ -272,6 +284,9 @@ export default function (pi) {
   }
 
   pi.events.on("herdr:blocked", (data) => {
+    if (!rootSession) {
+      return;
+    }
     if (!data?.active) {
       blockedCount = Math.max(0, blockedCount - 1);
       if (blockedCount === 0) {
@@ -287,7 +302,20 @@ export default function (pi) {
     publishState();
   });
 
+  pi.on("session_start", (_event, ctx) => {
+    if (ctx?.hasUI !== true) {
+      return;
+    }
+    rootSession = true;
+    updateSessionRef(ctx);
+    void reportSession();
+    publishState(true);
+  });
+
   pi.on("agent_start", () => {
+    if (!rootSession) {
+      return;
+    }
     clearPendingTimers();
     clearFailureState();
     agentActive = true;
@@ -295,6 +323,9 @@ export default function (pi) {
   });
 
   pi.on("agent_end", (event) => {
+    if (!rootSession) {
+      return;
+    }
     if (!agentActive) {
       // Pi can emit duplicate/late end events while auto-retry is already
       // holding the pane in Working. Do not let an unqualified duplicate end
@@ -312,9 +343,4 @@ export default function (pi) {
 
     scheduleIdle();
   });
-
-  pi.on("session_shutdown", async () => {
-    clearPendingTimers();
-    await releaseAgent();
-  });
 }