fix: preserve pi and omp working status across reloads (#984)

3 files changed, 136 insertions, 0 deletions

diff --git a/src/integration/assets/herdr-agent-state.test.ts b/src/integration/assets/herdr-agent-state.test.ts
new file mode 100644
index 00000000..652685cb
--- /dev/null
+++ b/src/integration/assets/herdr-agent-state.test.ts
@@ -0,0 +1,132 @@
+import { afterEach, expect, test } from "bun:test";
+import { rm } from "node:fs/promises";
+import { createServer, type Server } from "node:net";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+
+const originalEnvironment = {
+  HERDR_ENV: process.env.HERDR_ENV,
+  HERDR_PANE_ID: process.env.HERDR_PANE_ID,
+  HERDR_SOCKET_PATH: process.env.HERDR_SOCKET_PATH,
+};
+
+let server: Server | undefined;
+let socketPath: string | undefined;
+
+afterEach(async () => {
+  await new Promise<void>((resolve, reject) => {
+    if (!server) {
+      resolve();
+      return;
+    }
+    server.close((error) => (error ? reject(error) : resolve()));
+  });
+  server = undefined;
+
+  if (socketPath) {
+    await rm(socketPath, { force: true });
+    socketPath = undefined;
+  }
+
+  for (const [name, value] of Object.entries(originalEnvironment)) {
+    if (value === undefined) {
+      delete process.env[name];
+    } else {
+      process.env[name] = value;
+    }
+  }
+});
+
+const integrations = [
+  { name: "Pi", modulePath: "./pi/herdr-agent-state.ts" },
+  { name: "Oh My Pi", modulePath: "./omp/herdr-agent-state.ts" },
+] as const;
+
+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 });
+
+    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 = {
+      on(event: string, handler: Handler) {
+        handlers.set(event, handler);
+      },
+      events: {
+        on() {
+          return () => {};
+        },
+      },
+    };
+
+    const { default: install } = await import(integration.modulePath);
+    install(pi);
+
+    const sessionStart = handlers.get("session_start");
+    expect(sessionStart).toBeDefined();
+    await sessionStart?.(
+      { reason: "reload" },
+      {
+        hasUI: true,
+        isIdle: () => false,
+        sessionManager: {
+          getSessionFile: () => undefined,
+          getSessionId: () => undefined,
+        },
+      },
+    );
+
+    const reportedState = () => {
+      for (const request of requests) {
+        if (!isRecord(request) || request.method !== "pane.report_agent") {
+          continue;
+        }
+        const params = request.params;
+        if (isRecord(params) && typeof params.state === "string") {
+          return params.state;
+        }
+      }
+      return undefined;
+    };
+
+    const deadline = Date.now() + 1_000;
+    while (Date.now() < deadline && reportedState() === undefined) {
+      await Bun.sleep(5);
+    }
+
+    expect(reportedState()).toBe("working");
+  });
+}
+
+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 21c218b7..1b760a90 100644
--- a/src/integration/assets/omp/herdr-agent-state.ts
+++ b/src/integration/assets/omp/herdr-agent-state.ts
@@ -375,6 +375,8 @@ export default function (pi) {
     if (!activateRootSession(ctx)) {
       return;
     }
+    // A reload can replace this extension mid-run without emitting another agent_start.
+    agentActive = ctx?.isIdle?.() === false;
     publishState(true);
   });
 
diff --git a/src/integration/assets/pi/herdr-agent-state.ts b/src/integration/assets/pi/herdr-agent-state.ts
index 4bee2763..44a56088 100644
--- a/src/integration/assets/pi/herdr-agent-state.ts
+++ b/src/integration/assets/pi/herdr-agent-state.ts
@@ -332,6 +332,8 @@ export default function (pi) {
     rootSession = true;
     updateSessionRef(ctx);
     void reportSession();
+    // A reload can replace this extension mid-run without emitting another agent_start.
+    agentActive = ctx?.isIdle?.() === false;
     publishState(true);
   });