From ebff34063a55446c7390e12ec84acd2dad04ce43 Mon Sep 17 00:00:00 2001
From: Ogulcan Celik <ogulcancelik@gmail.com>
Date: Tue, 16 Jun 2026 17:57:24 +0300
Subject: fix: restore omp sessions after same-pane restart

refs #614
---
 src/integration/assets/omp/herdr-agent-state.ts | 44 +++++++++++++++++++++----
 src/integration/mod.rs                          |  5 ++-
 2 files changed, 41 insertions(+), 8 deletions(-)

(limited to 'src/integration')

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"));
-- 
cgit v1.2.3