fix: protect root agent restore sessions

refs #712

1 files changed, 32 insertions, 23 deletions

diff --git a/src/integration/assets/omp/herdr-agent-state.ts b/src/integration/assets/omp/herdr-agent-state.ts
index a6428cdd..a933e867 100644
--- a/src/integration/assets/omp/herdr-agent-state.ts
+++ b/src/integration/assets/omp/herdr-agent-state.ts
@@ -100,6 +100,35 @@ function parseDurationEnv(name: string, fallback: number): number {
   return parsed;
 }
 
+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: "omp",
+      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: "omp",
-      seq: nextReportSeq(),
-    },
-  });
-}
-
 export default function (pi) {
   if (!enabled()) {
     return;
@@ -287,11 +303,12 @@ export default function (pi) {
   });
 
   pi.on("session_start", (_event, ctx) => {
-    rootSession = ctx?.hasUI === true;
-    if (!rootSession) {
+    if (ctx?.hasUI !== true) {
       return;
     }
+    rootSession = true;
     updateSessionRef(ctx);
+    void reportSession();
     publishState(true);
   });
 
@@ -326,12 +343,4 @@ export default function (pi) {
 
     scheduleIdle();
   });
-
-  pi.on("session_shutdown", async () => {
-    if (!rootSession) {
-      return;
-    }
-    clearPendingTimers();
-    await releaseAgent();
-  });
 }