From 6a905814b59ade5055eb5fab5c96cfa596263dca Mon Sep 17 00:00:00 2001
From: Adam Cheng <63501289+627150795@users.noreply.github.com>
Date: Sun, 6 Sep 2026 14:11:44 +0800
Subject: [PATCH] fix(git-info): stop stale polling after session replacement

---
 extensions/git-info/index.ts            | 35 ++++++++-----
 tests/extensions/git-info/index.test.ts | 69 +++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 12 deletions(-)

diff --git a/extensions/git-info/index.ts b/extensions/git-info/index.ts
index b0d514b9..0b77b0d6 100644
--- a/extensions/git-info/index.ts
+++ b/extensions/git-info/index.ts
@@ -158,18 +158,24 @@ export default function gitInfo(pi: ExtensionAPI) {
   const refresh = (ctx: ExtensionContext, forcePullRequest = false) =>
     refreshCoordinator.run(refreshEffect(ctx, forcePullRequest, generation));
 
-  const refreshIfIdle = (ctx: ExtensionContext) =>
-    refreshCoordinator.runIfIdle(refreshEffect(ctx, false, generation));
+  const refreshIfIdle = (
+    ctx: ExtensionContext,
+    refreshGeneration = generation,
+  ) =>
+    refreshCoordinator.runIfIdle(refreshEffect(ctx, false, refreshGeneration));
 
   const reportBackgroundDefect = (defect: unknown) =>
     Effect.logError("git-info background task defect", defect);
 
-  const poll = () =>
+  const poll = (ctx: ExtensionContext, pollGeneration: number) =>
     Effect.suspend(() =>
-      currentContext ? refreshIfIdle(currentContext) : Effect.void,
+      pollGeneration === generation
+        ? refreshIfIdle(ctx, pollGeneration)
+        : Effect.fail("stale polling session"),
     ).pipe(
       Effect.catchDefect(reportBackgroundDefect),
       Effect.repeat(Schedule.fixed(POLL_INTERVAL_MS)),
+      Effect.catch(() => Effect.void),
       Effect.delay(POLL_INTERVAL_MS),
       Effect.asVoid,
     );
@@ -183,24 +189,29 @@ export default function gitInfo(pi: ExtensionAPI) {
     forkBackground(refreshIfIdle(ctx));
   };
 
+  const stopPolling = async (activeRuntime = runtime) => {
+    const previousPollingFiber = pollingFiber;
+    pollingFiber = undefined;
+    if (previousPollingFiber && activeRuntime) {
+      await activeRuntime.runPromise(Fiber.interrupt(previousPollingFiber));
+    }
+  };
+
   const stopRefreshListener = pi.events.on(REFRESH_CHANNEL, () => {
     if (currentContext) refreshInBackground(currentContext);
   });
 
   pi.on("session_start", async (_event, ctx) => {
-    generation += 1;
+    const sessionGeneration = ++generation;
     queriedPrBranch = null;
 
-    const previousPollingFiber = pollingFiber;
-    pollingFiber = undefined;
-    if (previousPollingFiber) {
-      await getRuntime().runPromise(Fiber.interrupt(previousPollingFiber));
-    }
+    await stopPolling();
+    if (sessionGeneration !== generation) return;
 
     // Do not block Pi startup on GitHub/network I/O. The initial refresh publishes
     // state when it completes; polling continues to keep it current afterwards.
     refreshInBackground(ctx);
-    pollingFiber = forkBackground(poll());
+    pollingFiber = forkBackground(poll(ctx, sessionGeneration));
   });
 
   pi.on("input", (_event, ctx) => {
@@ -216,8 +227,8 @@ export default function gitInfo(pi: ExtensionAPI) {
     stopRefreshListener();
     generation += 1;
     currentContext = undefined;
-    pollingFiber = undefined;
     const closing = runtime;
+    await stopPolling(closing);
     runtime = undefined;
     await closing?.dispose();
   });
diff --git a/tests/extensions/git-info/index.test.ts b/tests/extensions/git-info/index.test.ts
index efc068e9..bc8061fc 100644
--- a/tests/extensions/git-info/index.test.ts
+++ b/tests/extensions/git-info/index.test.ts
@@ -186,3 +186,72 @@ printf '%s\\n' '{"number":42,"url":"https://example.test/pr/42","state":"OPEN","
     rmSync(root, { recursive: true, force: true });
   }
 });
+
+test("rapid session starts do not leave a stale polling context", async () => {
+  const root = process.cwd();
+  const hooks = new Map<
+    string,
+    (event: unknown, ctx: ExtensionContext) => unknown
+  >();
+  let activeSession = 1;
+  let staleReads = 0;
+  let currentSessionReads = 0;
+  const api = {
+    events: {
+      on: () => () => undefined,
+      emit: () => undefined,
+    },
+    on: (
+      event: string,
+      handler: (event: unknown, ctx: ExtensionContext) => unknown,
+    ) => {
+      hooks.set(event, handler);
+    },
+    registerCommand: () => undefined,
+  } as unknown as ExtensionAPI;
+  const makeContext = (session: number) => {
+    const ctx = {
+      mode: "tui",
+      signal: undefined,
+      ui: { notify: () => undefined },
+    } as Record<string, unknown>;
+    Object.defineProperty(ctx, "cwd", {
+      get: () => {
+        if (session !== activeSession) {
+          staleReads += 1;
+          throw new Error(
+            "This extension ctx is stale after session replacement",
+          );
+        }
+        currentSessionReads += 1;
+        return root;
+      },
+    });
+    return ctx as unknown as ExtensionContext;
+  };
+  const start = (event: unknown, ctx: ExtensionContext) =>
+    hooks.get("session_start")?.(event, ctx);
+  const ctx1 = makeContext(1);
+
+  gitInfo(api);
+  try {
+    await start({}, ctx1);
+    await new Promise((resolve) => setTimeout(resolve, 100));
+
+    activeSession = 10;
+    const starts = Array.from({ length: 9 }, (_, index) =>
+      start({}, makeContext(index + 2)),
+    );
+    await Promise.all(starts);
+    await new Promise((resolve) => setTimeout(resolve, 1_000));
+    assert.equal(currentSessionReads > 0, true);
+    currentSessionReads = 0;
+    staleReads = 0;
+    await new Promise((resolve) => setTimeout(resolve, 6_000));
+
+    assert.equal(currentSessionReads > 0, true);
+    assert.equal(staleReads, 0);
+  } finally {
+    await hooks.get("session_shutdown")?.({}, makeContext(activeSession));
+  }
+});