2 files changed, 166 insertions, 56 deletions
diff --git a/src/integration/assets/herdr-agent-state.test.ts b/src/integration/assets/herdr-agent-state.test.ts
index 69cb8894..077d1894 100644
--- a/src/integration/assets/herdr-agent-state.test.ts
+++ b/src/integration/assets/herdr-agent-state.test.ts
@@ -48,42 +48,13 @@ function importFresh(modulePath: string) {
   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(
-      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);
-    });
+type Handler = (event: unknown, context: unknown) => unknown;
 
-    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 = {
+function createExtensionHarness() {
+  const handlers = new Map<string, Handler>();
+  return {
+    handlers,
+    pi: {
       on(event: string, handler: Handler) {
         handlers.set(event, handler);
       },
@@ -92,7 +63,50 @@ for (const integration of integrations) {
           return () => {};
         },
       },
-    };
+    },
+  };
+}
+
+function configureIntegrationEnvironment(recordingSocketPath: string) {
+  process.env.HERDR_ENV = "1";
+  process.env.HERDR_SOCKET_PATH = recordingSocketPath;
+  process.env.HERDR_PANE_ID = "test:p1";
+}
+
+async function startRecordingServer(name: string): Promise<unknown[]> {
+  const recordingSocketPath = join(tmpdir(), `herdr-${name}-${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);
+  });
+  configureIntegrationEnvironment(recordingSocketPath);
+  return requests;
+}
+
+for (const integration of integrations) {
+  test(`${integration.name} reload preserves working state when the agent is active`, async () => {
+    const requests = await startRecordingServer(
+      integration.name.toLowerCase().replaceAll(" ", "-"),
+    );
+    const { handlers, pi } = createExtensionHarness();
 
     const { default: install } = await importFresh(integration.modulePath);
     install(pi);
@@ -133,6 +147,115 @@ for (const integration of integrations) {
   });
 }
 
+test("Pi reports the session replacement source", async () => {
+  const requests = await startRecordingServer("pi-session-source");
+  const { handlers, pi } = createExtensionHarness();
+
+  const { default: install } = await importFresh("./pi/herdr-agent-state.ts");
+  install(pi);
+
+  const sessionStart = handlers.get("session_start");
+  expect(sessionStart).toBeDefined();
+  await sessionStart?.(
+    { reason: "new" },
+    {
+      hasUI: true,
+      isIdle: () => true,
+      sessionManager: {
+        getSessionFile: () => "/tmp/pi-new.jsonl",
+        getSessionId: () => "pi-new",
+      },
+    },
+  );
+
+  const reportedSession = () =>
+    requests.find((request) => isRecord(request) && request.method === "pane.report_agent_session");
+  const deadline = Date.now() + 1_000;
+  while (Date.now() < deadline && reportedSession() === undefined) {
+    await Bun.sleep(5);
+  }
+
+  const request = reportedSession();
+  expect(request).toBeDefined();
+  expect(isRecord(request) && isRecord(request.params) ? request.params.session_start_source : null)
+    .toBe("new");
+});
+
+test("Pi waits for a replacement session report before publishing state", async () => {
+  const recordingSocketPath = join(tmpdir(), `herdr-pi-session-order-${process.pid}.sock`);
+  socketPath = recordingSocketPath;
+  await rm(recordingSocketPath, { force: true });
+
+  const requests: unknown[] = [];
+  let acknowledgeSessionReport: (() => void) | undefined;
+  const recordingServer = createServer((socket) => {
+    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));
+      requests.push(request);
+      if (isRecord(request) && request.method === "pane.report_agent_session") {
+        acknowledgeSessionReport = () => socket.end("{}\n");
+        return;
+      }
+      socket.end("{}\n");
+    });
+  });
+  server = recordingServer;
+  await new Promise<void>((resolve, reject) => {
+    recordingServer.once("error", reject);
+    recordingServer.listen(recordingSocketPath, resolve);
+  });
+
+  configureIntegrationEnvironment(recordingSocketPath);
+  const { handlers, pi } = createExtensionHarness();
+  const { default: install } = await importFresh("./pi/herdr-agent-state.ts");
+  install(pi);
+
+  const sessionStart = handlers.get("session_start");
+  expect(sessionStart).toBeDefined();
+  const sessionStartResult = sessionStart?.(
+    { reason: "new" },
+    {
+      hasUI: true,
+      isIdle: () => false,
+      sessionManager: {
+        getSessionFile: () => "/tmp/pi-new.jsonl",
+        getSessionId: () => "pi-new",
+      },
+    },
+  );
+
+  const deadline = Date.now() + 1_000;
+  while (Date.now() < deadline && acknowledgeSessionReport === undefined) {
+    await Bun.sleep(5);
+  }
+  expect(acknowledgeSessionReport).toBeDefined();
+  expect(
+    requests.some((request) => isRecord(request) && request.method === "pane.report_agent"),
+  ).toBe(false);
+
+  acknowledgeSessionReport?.();
+  await sessionStartResult;
+
+  const stateDeadline = Date.now() + 1_000;
+  while (
+    Date.now() < stateDeadline &&
+    !requests.some((request) => isRecord(request) && request.method === "pane.report_agent")
+  ) {
+    await Bun.sleep(5);
+  }
+  expect(requests.map((request) => (isRecord(request) ? request.method : undefined))).toEqual([
+    "pane.report_agent_session",
+    "pane.report_agent",
+  ]);
+});
+
 test("Pi retries working state after an unanswered socket attempt", async () => {
   const recordingSocketPath = join(tmpdir(), `herdr-pi-retry-${process.pid}.sock`);
   socketPath = recordingSocketPath;
@@ -167,22 +290,8 @@ test("Pi retries working state after an unanswered socket attempt", async () =>
     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 () => {};
-      },
-    },
-  };
+  configureIntegrationEnvironment(recordingSocketPath);
+  const { handlers, pi } = createExtensionHarness();
 
   const { default: install } = await importFresh("./pi/herdr-agent-state.ts");
   install(pi);
diff --git a/src/integration/assets/pi/herdr-agent-state.ts b/src/integration/assets/pi/herdr-agent-state.ts
index 9048911f..52340ebb 100644
--- a/src/integration/assets/pi/herdr-agent-state.ts
+++ b/src/integration/assets/pi/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=pi
-// HERDR_INTEGRATION_VERSION=4
+// HERDR_INTEGRATION_VERSION=5
 // @ts-nocheck
 
 import { createConnection } from "node:net";
@@ -121,7 +121,7 @@ function currentSessionRef(): Record<string, unknown> | undefined {
   return undefined;
 }
 
-function reportSession(): Promise<void> {
+function reportSession(sessionStartSource?: string): Promise<void> {
   const sessionRef = currentSessionRef();
   if (!sessionRef) {
     return Promise.resolve();
@@ -135,6 +135,7 @@ function reportSession(): Promise<void> {
       source,
       agent: "pi",
       seq: nextReportSeq(),
+      session_start_source: sessionStartSource,
       ...sessionRef,
     },
   });
@@ -336,13 +337,13 @@ export default function (pi) {
     publishState();
   });
 
-  pi.on("session_start", (_event, ctx) => {
+  pi.on("session_start", async (event, ctx) => {
     if (ctx?.hasUI !== true) {
       return;
     }
     rootSession = true;
     updateSessionRef(ctx);
-    void reportSession();
+    await reportSession(event?.reason);
     // A reload can replace this extension mid-run without emitting another agent_start.
     agentActive = ctx?.isIdle?.() === false;
     publishState(true);