fix: retry dropped omp lifecycle reports

refs #1310

4 files changed, 102 insertions, 13 deletions

diff --git a/src/integration/assets/herdr-agent-state.test.ts b/src/integration/assets/herdr-agent-state.test.ts
index 077d1894..c662e132 100644
--- a/src/integration/assets/herdr-agent-state.test.ts
+++ b/src/integration/assets/herdr-agent-state.test.ts
@@ -6,6 +6,7 @@ import { join } from "node:path";
 
 const originalEnvironment = {
   HERDR_ENV: process.env.HERDR_ENV,
+  HERDR_OMP_IDLE_DEBOUNCE_MS: process.env.HERDR_OMP_IDLE_DEBOUNCE_MS,
   HERDR_PANE_ID: process.env.HERDR_PANE_ID,
   HERDR_SOCKET_PATH: process.env.HERDR_SOCKET_PATH,
 };
@@ -256,8 +257,8 @@ test("Pi waits for a replacement session report before publishing state", async
   ]);
 });
 
-test("Pi retries working state after an unanswered socket attempt", async () => {
-  const recordingSocketPath = join(tmpdir(), `herdr-pi-retry-${process.pid}.sock`);
+async function startDroppedFirstResponseServer(name: string) {
+  const recordingSocketPath = join(tmpdir(), `herdr-${name}-${process.pid}.sock`);
   socketPath = recordingSocketPath;
   await rm(recordingSocketPath, { force: true });
 
@@ -291,6 +292,46 @@ test("Pi retries working state after an unanswered socket attempt", async () =>
   });
 
   configureIntegrationEnvironment(recordingSocketPath);
+  return {
+    attemptedRequests,
+    deliveredRequests,
+    connectionCount: () => connectionCount,
+  };
+}
+
+test("Oh My Pi retries working before a queued idle state", async () => {
+  const { attemptedRequests } = await startDroppedFirstResponseServer("omp-retry");
+  process.env.HERDR_OMP_IDLE_DEBOUNCE_MS = "0";
+  const { handlers, pi } = createExtensionHarness();
+
+  const { default: install } = await importFresh("./omp/herdr-agent-state.ts");
+  install(pi);
+
+  const context = {
+    hasUI: true,
+    isIdle: () => false,
+    sessionManager: {
+      getSessionFile: () => undefined,
+      getSessionId: () => undefined,
+    },
+  };
+  handlers.get("session_start")?.({ reason: "startup" }, context);
+  handlers.get("agent_end")?.({ messages: [] }, context);
+
+  const deadline = Date.now() + 2_500;
+  while (Date.now() < deadline && attemptedRequests.length < 3) {
+    await Bun.sleep(5);
+  }
+
+  expect(attemptedRequests).toHaveLength(3);
+  expect(attemptedRequests[1]).toEqual(attemptedRequests[0]);
+  expect(requestState(attemptedRequests[0])).toBe("working");
+  expect(requestState(attemptedRequests[2])).toBe("idle");
+});
+
+test("Pi retries working state after an unanswered socket attempt", async () => {
+  const { attemptedRequests, deliveredRequests, connectionCount } =
+    await startDroppedFirstResponseServer("pi-retry");
   const { handlers, pi } = createExtensionHarness();
 
   const { default: install } = await importFresh("./pi/herdr-agent-state.ts");
@@ -324,12 +365,19 @@ test("Pi retries working state after an unanswered socket attempt", async () =>
     await Bun.sleep(5);
   }
 
-  expect(connectionCount).toBeGreaterThanOrEqual(2);
+  expect(connectionCount()).toBeGreaterThanOrEqual(2);
   expect(attemptedRequests.length).toBeGreaterThanOrEqual(2);
   expect(attemptedRequests[1]).toEqual(attemptedRequests[0]);
   expect(reportedWorking()).toBe(true);
 });
 
+function requestState(request: unknown): unknown {
+  if (!isRecord(request) || !isRecord(request.params)) {
+    return undefined;
+  }
+  return request.params.state;
+}
+
 function isRecord(value: unknown): value is Record<string, unknown> {
   return typeof value === "object" && value !== null;
 }
diff --git a/src/integration/assets/omp/herdr-agent-state.ts b/src/integration/assets/omp/herdr-agent-state.ts
index 1b760a90..9bb6671a 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=4
+// HERDR_INTEGRATION_VERSION=5
 // @ts-nocheck
 
 import { createConnection } from "node:net";
@@ -18,30 +18,41 @@ function enabled() {
 
 let requestQueue = Promise.resolve();
 
-function sendRequestNow(request: unknown): Promise<void> {
+function sendRequestAttempt(request: unknown, timeoutMs: number): Promise<boolean> {
   if (!enabled()) {
-    return Promise.resolve();
+    return Promise.resolve(true);
   }
 
   return new Promise((resolve) => {
     let done = false;
-    const finish = () => {
+    let timeout: ReturnType<typeof setTimeout> | undefined;
+    const finish = (delivered: boolean) => {
       if (done) return;
       done = true;
+      if (timeout) {
+        clearTimeout(timeout);
+      }
       socket.destroy();
-      resolve();
+      resolve(delivered);
     };
 
     const socket = createConnection(socketPath!);
-    socket.on("error", finish);
+    socket.on("error", () => finish(false));
     socket.on("connect", () => socket.write(`${JSON.stringify(request)}\n`));
-    socket.on("data", finish);
-    socket.on("end", finish);
-    const timeout = setTimeout(finish, 500);
+    socket.on("data", () => finish(true));
+    socket.on("end", () => finish(false));
+    timeout = setTimeout(() => finish(false), timeoutMs);
     timeout.unref?.();
   });
 }
 
+async function sendRequestNow(request: unknown): Promise<void> {
+  if (await sendRequestAttempt(request, 500)) {
+    return;
+  }
+  await sendRequestAttempt(request, 1500);
+}
+
 function sendRequest(request: unknown): Promise<void> {
   requestQueue = requestQueue.then(
     () => sendRequestNow(request),
diff --git a/src/integration/mod.rs b/src/integration/mod.rs
index 39414cbf..9cf3d24b 100644
--- a/src/integration/mod.rs
+++ b/src/integration/mod.rs
@@ -25,7 +25,7 @@ const PI_EXTENSION_ASSET: &str = include_str!("assets/pi/herdr-agent-state.ts");
 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;
+const OMP_INTEGRATION_VERSION: u32 = 5;
 const CLAUDE_HOOK_INSTALL_NAME: &str = if cfg!(windows) {
     "herdr-agent-state.ps1"
 } else {
diff --git a/src/integration/tests.rs b/src/integration/tests.rs
index 5d8538c7..3d578c53 100644
--- a/src/integration/tests.rs
+++ b/src/integration/tests.rs
@@ -763,6 +763,36 @@ fn outdated_integrations_detect_previous_pi_version() {
 }
 
 #[test]
+fn outdated_integrations_detect_previous_omp_version() {
+    let _lock = integration_env_lock();
+    let base = unique_base();
+    let home = base.join("home");
+    let ext_dir = home.join(".omp/agent/extensions");
+    fs::create_dir_all(&ext_dir).unwrap();
+    let extension_path = ext_dir.join(OMP_EXTENSION_INSTALL_NAME);
+    fs::write(
+        &extension_path,
+        "// HERDR_INTEGRATION_ID=omp\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::Omp
+    );
+    assert_eq!(outdated[0].path, extension_path);
+    assert_eq!(outdated[0].installed_version, Some(4));
+    assert_eq!(outdated[0].expected_version, OMP_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();
     let base = unique_base();