fix: retry pi state socket reports

refs #1049

2 files changed, 114 insertions, 9 deletions

diff --git a/src/integration/assets/herdr-agent-state.test.ts b/src/integration/assets/herdr-agent-state.test.ts
index 652685cb..69cb8894 100644
--- a/src/integration/assets/herdr-agent-state.test.ts
+++ b/src/integration/assets/herdr-agent-state.test.ts
@@ -12,6 +12,7 @@ const originalEnvironment = {
 
 let server: Server | undefined;
 let socketPath: string | undefined;
+let importCounter = 0;
 
 afterEach(async () => {
   await new Promise<void>((resolve, reject) => {
@@ -42,6 +43,11 @@ const integrations = [
   { name: "Oh My Pi", modulePath: "./omp/herdr-agent-state.ts" },
 ] as const;
 
+function importFresh(modulePath: string) {
+  importCounter += 1;
+  return import(`${modulePath}?test=${importCounter}`);
+}
+
 for (const integration of integrations) {
   test(`${integration.name} reload preserves working state when the agent is active`, async () => {
     const recordingSocketPath = join(
@@ -88,7 +94,7 @@ for (const integration of integrations) {
       },
     };
 
-    const { default: install } = await import(integration.modulePath);
+    const { default: install } = await importFresh(integration.modulePath);
     install(pi);
 
     const sessionStart = handlers.get("session_start");
@@ -127,6 +133,94 @@ for (const integration of integrations) {
   });
 }
 
+test("Pi retries working state after an unanswered socket attempt", async () => {
+  const recordingSocketPath = join(tmpdir(), `herdr-pi-retry-${process.pid}.sock`);
+  socketPath = recordingSocketPath;
+  await rm(recordingSocketPath, { force: true });
+
+  let connectionCount = 0;
+  const attemptedRequests: unknown[] = [];
+  const deliveredRequests: unknown[] = [];
+  const recordingServer = createServer((socket) => {
+    connectionCount += 1;
+    const connectionNumber = connectionCount;
+    let input = "";
+    socket.setEncoding("utf8");
+    socket.on("data", (chunk) => {
+      input += chunk;
+      const newline = input.indexOf("\n");
+      if (newline === -1) {
+        return;
+      }
+      const request = JSON.parse(input.slice(0, newline));
+      attemptedRequests.push(request);
+      if (connectionNumber === 1) {
+        return;
+      }
+      deliveredRequests.push(request);
+      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 importFresh("./pi/herdr-agent-state.ts");
+  install(pi);
+
+  const sessionStart = handlers.get("session_start");
+  expect(sessionStart).toBeDefined();
+  await sessionStart?.(
+    { reason: "startup" },
+    {
+      hasUI: true,
+      isIdle: () => false,
+      sessionManager: {
+        getSessionFile: () => undefined,
+        getSessionId: () => undefined,
+      },
+    },
+  );
+
+  const reportedWorking = () =>
+    deliveredRequests.some((request) => {
+      if (!isRecord(request) || request.method !== "pane.report_agent") {
+        return false;
+      }
+      const params = request.params;
+      return isRecord(params) && params.state === "working";
+    });
+
+  const deadline = Date.now() + 2_500;
+  while (Date.now() < deadline && !reportedWorking()) {
+    await Bun.sleep(5);
+  }
+
+  expect(connectionCount).toBeGreaterThanOrEqual(2);
+  expect(attemptedRequests.length).toBeGreaterThanOrEqual(2);
+  expect(attemptedRequests[1]).toEqual(attemptedRequests[0]);
+  expect(reportedWorking()).toBe(true);
+});
+
 function isRecord(value: unknown): value is Record<string, unknown> {
   return typeof value === "object" && value !== null;
 }
diff --git a/src/integration/assets/pi/herdr-agent-state.ts b/src/integration/assets/pi/herdr-agent-state.ts
index 44a56088..9048911f 100644
--- a/src/integration/assets/pi/herdr-agent-state.ts
+++ b/src/integration/assets/pi/herdr-agent-state.ts
@@ -16,30 +16,41 @@ function enabled() {
   return HERDR_ENV === "1" && !!socketPath && !!paneId;
 }
 
-function sendRequest(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 sendRequest(request: unknown): Promise<void> {
+  if (await sendRequestAttempt(request, 500)) {
+    return;
+  }
+  await sendRequestAttempt(request, 1500);
+}
+
 type AgentState = "working" | "blocked" | "idle";
 
 type QueuedState = {