Merge pull request #300 from kevin9327/fix/overlay-state-generation

fix(extension): ignore stale overlay state by generation

apps/extension/src/entrypoints/__tests__/background-overlays.test.ts

@@ -97,14 +97,20 @@ it.each([false, true])("hides an observed tab after detachment, remote=%s", asyn
   f.task.observedTabs = new Set([20]);
   f.deps.onAgentTabClaimed?.(20, 10);
   expect(shouldShowAgentControlOverlay(f.overlay.snapshot())).toBe(true);
+  const claimed = f.sendMessage.mock.lastCall![1];
   f.onDetached.emit(20);
   expect(isAgentControlledTab(f.task, 20)).toBe(false);
-  expect(f.sendMessage).toHaveBeenLastCalledWith(20, {
+  const [tabId, released] = f.sendMessage.mock.lastCall!;
+  expect(tabId).toBe(20);
+  expect(released).toMatchObject({
     type: OVERLAY_AGENT_STATE,
     sessionId: null,
     mode: "hidden",
-    generation: 0,
   });
+  // The release is ordered after the claim, although no session changed.
+  expect(released.epoch).toEqual(expect.any(String));
+  expect(released.epoch).toBe(claimed.epoch);
+  expect(released.generation).toBeGreaterThan(claimed.generation);
   expect(shouldShowAgentControlOverlay(f.overlay.snapshot())).toBe(false);
   expect(f.deps.cdp?.releaseSessionTab).toHaveBeenCalledWith("one", 20);
 });

apps/extension/src/entrypoints/__tests__/content-passthrough.test.ts

@@ -1,9 +1,10 @@
-import { afterEach, expect, it, vi } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
 import { INPUT_PASSTHROUGH, INPUT_PASSTHROUGH_ATTR } from "@/lib/input-passthrough-bridge";
 import {
   OVERLAY_AGENT_OVERLAY_RESET,
   OVERLAY_AGENT_STATE,
   type OverlayMode,
+  type OverlayVersion,
 } from "@/lib/overlay-bridge";
 import { RECORD_START, RECORD_STOP } from "@/lib/record-bridge";
 
@@ -81,12 +82,21 @@ async function fixture() {
     },
   });
   const send = (message: unknown) => receive(message, {}, vi.fn());
-  const state = (sessionId = "one", mode: OverlayMode = "control") =>
-    send({ type: OVERLAY_AGENT_STATE, sessionId, mode, generation: 0 });
+  // Versions in the order one background worker hands them out.
+  let generation = 0;
+  const version = (): OverlayVersion => ({ epoch: "worker", generation: ++generation });
+  const state = (
+    sessionId: string | null = "one",
+    mode: OverlayMode = "control",
+    v: OverlayVersion = version(),
+  ) => send({ type: OVERLAY_AGENT_STATE, sessionId, mode, ...v });
+  const reset = (sessionId: string, v: OverlayVersion = version()) =>
+    send({ type: OVERLAY_AGENT_OVERLAY_RESET, sessionId, ...v });
   const begin = (id = "click-one") => send({ type: INPUT_PASSTHROUGH, phase: "begin", id });
   const active = () => host.hasAttribute(INPUT_PASSTHROUGH_ATTR);
+  const blocking = () => host.hasAttribute("data-bsk-overlay-blocking");
   state();
-  return { send, state, begin, active, remount: () => mount() };
+  return { send, version, state, reset, begin, active, blocking, remount: () => mount() };
 }
 
 it.each([
@@ -121,10 +131,10 @@ it("clears old session leases without letting stale reset/end clear the new sess
   f.state("two");
   expect(f.active()).toBe(false);
   f.begin("click-two");
-  f.send({ type: OVERLAY_AGENT_OVERLAY_RESET, sessionId: "one" });
+  f.reset("one");
   f.send({ type: INPUT_PASSTHROUGH, phase: "end", id: "click-one" });
   expect(f.active()).toBe(true);
-  f.send({ type: OVERLAY_AGENT_OVERLAY_RESET, sessionId: "two" });
+  f.reset("two");
   expect(f.active()).toBe(false);
   expect(vi.getTimerCount()).toBe(0);
   f.remount();
@@ -149,3 +159,54 @@ it("removes passthrough and timers when the content context is invalidated", asy
   expect(f.active()).toBe(false);
   expect(vi.getTimerCount()).toBe(0);
 });
+
+describe("overlay state ordering", () => {
+  it("drops a control state that arrives after a newer hidden state", async () => {
+    const f = await fixture();
+    const late = f.version();
+    f.state(null, "hidden");
+    f.state("one", "control", late);
+    expect(f.blocking()).toBe(false);
+  });
+
+  it("drops a control state sent before a reset", async () => {
+    const f = await fixture();
+    expect(f.blocking()).toBe(true);
+    const late = f.version();
+    f.reset("one");
+    expect(f.blocking()).toBe(false);
+    f.state("one", "control", late);
+    expect(f.blocking()).toBe(false);
+  });
+
+  it("drops a reset sent before a newer control state", async () => {
+    const f = await fixture();
+    const late = f.version();
+    f.state("one", "control");
+    f.reset("one", late);
+    expect(f.blocking()).toBe(true);
+  });
+
+  it("follows a restarted worker although its generation starts over", async () => {
+    const f = await fixture();
+    f.state("one", "control", { epoch: "worker", generation: 20 });
+    f.state(null, "hidden", { epoch: "restarted", generation: 1 });
+    expect(f.blocking()).toBe(false);
+    f.state("one", "control", { epoch: "restarted", generation: 2 });
+    expect(f.blocking()).toBe(true);
+  });
+
+  it("lets a restarted worker claim a page it had hidden", async () => {
+    const f = await fixture();
+    f.state(null, "hidden", { epoch: "worker", generation: 20 });
+    f.state("one", "control", { epoch: "restarted", generation: 1 });
+    expect(f.blocking()).toBe(true);
+  });
+
+  it("keeps dropping the replaced worker's messages", async () => {
+    const f = await fixture();
+    f.state(null, "hidden", { epoch: "restarted", generation: 1 });
+    f.state("one", "control", { epoch: "worker", generation: 30 });
+    expect(f.blocking()).toBe(false);
+  });
+});

apps/extension/src/entrypoints/background.ts

@@ -27,6 +27,7 @@ import {
   type OverlayMessage,
   type OverlayMode,
 } from "@/lib/overlay-bridge";
+import { nextOverlayVersion } from "@/lib/overlay-version";
 import { POPUP_PORT_NAME, type PopupInbound, type PopupOutbound } from "@/lib/popup-bridge";
 import { recordFrameCoordinator } from "@/lib/recording/frame-coordinator";
 import { attachSessionsLiveFlag } from "@/lib/sessions-live-flag";
@@ -99,7 +100,6 @@ export default defineBackground(() => {
   });
   const sessionsLive = attachSessionsLiveFlag({ manager: sessions });
   const popupSnapshotRefreshers = new Set<() => void>();
-  let overlayGeneration = 0;
   const controlModes = new Map<string, OverlayMode>();
 
   watchRemoteAuthorization();
@@ -140,7 +140,6 @@ export default defineBackground(() => {
   function setControlMode(sessionId: string, mode: OverlayMode): void {
     if (controlModes.get(sessionId) === mode) return;
     controlModes.set(sessionId, mode);
-    overlayGeneration += 1;
     const ctx = sessions.get(sessionId);
     if (ctx) void pushOverlayStateForWindow(ctx.agentWindowId);
   }
@@ -152,14 +151,14 @@ export default defineBackground(() => {
         type: OVERLAY_AGENT_STATE,
         sessionId: null,
         mode: "hidden",
-        generation: overlayGeneration,
+        ...nextOverlayVersion(),
       };
     }
     return {
       type: OVERLAY_AGENT_STATE,
       sessionId: ctx.sessionId,
       mode: controlModes.get(ctx.sessionId) ?? "control",
-      generation: overlayGeneration,
+      ...nextOverlayVersion(),
     };
   }
 
@@ -177,7 +176,7 @@ export default defineBackground(() => {
       type: OVERLAY_AGENT_STATE,
       sessionId: null,
       mode: "hidden",
-      generation: overlayGeneration,
+      ...nextOverlayVersion(),
     };
   }
 
@@ -221,7 +220,6 @@ export default defineBackground(() => {
     for (const sessionId of controlModes.keys()) {
       if (!liveSessionIds.has(sessionId)) controlModes.delete(sessionId);
     }
-    overlayGeneration += 1;
     pushAllAgentOverlayStates();
     for (const refresh of popupSnapshotRefreshers) refresh();
   }

apps/extension/src/entrypoints/content.ts

@@ -43,6 +43,7 @@ import {
   type OverlayAgentOverlayResetMessage,
   type OverlayAgentStateMessage,
   type OverlayAutomationBypassMessage,
+  OverlayVersionGate,
 } from "@/lib/overlay-bridge";
 import { sendInterrupt } from "@/lib/overlay-interrupt-client";
 import {
@@ -81,6 +82,7 @@ export default defineContentScript({
     let overlayHost: HTMLElement | null = null;
     let overlayContainer: HTMLElement | null = null;
     let activeAgentState: OverlayAgentStateMessage | null = null;
+    const overlayVersions = new OverlayVersionGate();
     let hostLossReported = false;
     let remountInProgress = false;
 
@@ -227,6 +229,10 @@ export default defineContentScript({
       resetAgentOverlayState(sessionId);
     }
 
+    function receiveOverlayState(state: OverlayAgentStateMessage): void {
+      if (overlayVersions.admit(state)) applyOverlayState(state);
+    }
+
     function applyOverlayState(state: OverlayAgentStateMessage): void {
       if (overlays.snapshot().activeSessionId !== state.sessionId) inputPassthrough.reset();
       activeAgentState = state;
@@ -331,12 +337,12 @@ export default defineContentScript({
       }
 
       if (isOverlayAgentStateMessage(message)) {
-        applyOverlayState(message);
+        receiveOverlayState(message);
         return false;
       }
 
       if (isOverlayAgentOverlayResetMessage(message)) {
-        resetAgentOverlayState(message.sessionId);
+        if (overlayVersions.admit(message)) resetAgentOverlayState(message.sessionId);
         return false;
       }
 
@@ -481,7 +487,7 @@ export default defineContentScript({
           kind: OVERLAY_MSG_READY,
         })) as OverlayAgentStateMessage | undefined;
         if (state && isOverlayAgentStateMessage(state)) {
-          applyOverlayState(state);
+          receiveOverlayState(state);
         }
         void refreshAuxiliaryOverlayState();
       } catch (err) {

apps/extension/src/lib/__tests__/overlay-bridge.test.ts

@@ -1,10 +1,13 @@
 import { describe, expect, it } from "vitest";
 import {
   isOverlayAgentOverlayResetMessage,
+  isOverlayAgentStateMessage,
   OVERLAY_AGENT_OVERLAY_RESET,
+  OVERLAY_AGENT_STATE,
   OVERLAY_MSG_INTERRUPT,
   type OverlayInterruptRequest,
   type OverlayInterruptResponse,
+  OverlayVersionGate,
 } from "@/lib/overlay-bridge";
 
 describe("OVERLAY_MSG_INTERRUPT", () => {
@@ -28,13 +31,24 @@ describe("OVERLAY_MSG_INTERRUPT", () => {
 });
 
 describe("isOverlayAgentOverlayResetMessage", () => {
-  it("accepts reset messages with a session id", () => {
+  it("accepts reset messages with a session id and a version", () => {
+    expect(
+      isOverlayAgentOverlayResetMessage({
+        type: OVERLAY_AGENT_OVERLAY_RESET,
+        sessionId: "sess-1",
+        epoch: "worker",
+        generation: 1,
+      }),
+    ).toBe(true);
+  });
+
+  it("rejects reset messages without a version", () => {
     expect(
       isOverlayAgentOverlayResetMessage({
         type: OVERLAY_AGENT_OVERLAY_RESET,
         sessionId: "sess-1",
       }),
-    ).toBe(true);
+    ).toBe(false);
   });
 
   it("rejects reset messages without a session id", () => {
@@ -45,3 +59,34 @@ describe("isOverlayAgentOverlayResetMessage", () => {
     ).toBe(false);
   });
 });
+
+describe("isOverlayAgentStateMessage", () => {
+  it("requires a version", () => {
+    const message = { type: OVERLAY_AGENT_STATE, sessionId: "sess-1", mode: "control" };
+    expect(isOverlayAgentStateMessage({ ...message, epoch: "worker", generation: 1 })).toBe(true);
+    expect(isOverlayAgentStateMessage({ ...message, generation: 1 })).toBe(false);
+  });
+});
+
+describe("OverlayVersionGate", () => {
+  it("admits the first message and the same message again", () => {
+    const gate = new OverlayVersionGate();
+    expect(gate.admit({ epoch: "worker", generation: 4 })).toBe(true);
+    expect(gate.admit({ epoch: "worker", generation: 4 })).toBe(true);
+  });
+
+  it("orders messages from one worker by generation", () => {
+    const gate = new OverlayVersionGate();
+    expect(gate.admit({ epoch: "worker", generation: 6 })).toBe(true);
+    expect(gate.admit({ epoch: "worker", generation: 5 })).toBe(false);
+    expect(gate.admit({ epoch: "worker", generation: 7 })).toBe(true);
+  });
+
+  it("follows a restarted worker and drops the one it replaced", () => {
+    const gate = new OverlayVersionGate();
+    expect(gate.admit({ epoch: "worker", generation: 20 })).toBe(true);
+    expect(gate.admit({ epoch: "restarted", generation: 1 })).toBe(true);
+    expect(gate.admit({ epoch: "worker", generation: 21 })).toBe(false);
+    expect(gate.admit({ epoch: "restarted", generation: 2 })).toBe(true);
+  });
+});

apps/extension/src/lib/__tests__/overlay-version.test.ts

@@ -0,0 +1,28 @@
+import { afterEach, expect, it, vi } from "vitest";
+import { OVERLAY_AGENT_OVERLAY_RESET } from "@/lib/overlay-bridge";
+import { nextOverlayVersion } from "@/lib/overlay-version";
+import { chromeAgentOverlayResetApi } from "@/tools/tabs";
+
+afterEach(() => {
+  vi.unstubAllGlobals();
+});
+
+it("hands out increasing generations under one epoch", () => {
+  const first = nextOverlayVersion();
+  const second = nextOverlayVersion();
+  expect(second.epoch).toBe(first.epoch);
+  expect(second.generation).toBeGreaterThan(first.generation);
+});
+
+it("versions a reset like a state, so the page can order the two", async () => {
+  const sendMessage = vi.fn(async () => {});
+  vi.stubGlobal("chrome", { tabs: { sendMessage } });
+  const before = nextOverlayVersion();
+  await chromeAgentOverlayResetApi.resetAgentOverlays(3, "sess-1");
+  expect(sendMessage).toHaveBeenCalledWith(3, {
+    type: OVERLAY_AGENT_OVERLAY_RESET,
+    sessionId: "sess-1",
+    epoch: before.epoch,
+    generation: before.generation + 1,
+  });
+});

apps/extension/src/lib/overlay-bridge.ts

@@ -68,11 +68,20 @@ export type OverlayMode = "control" | "interrupting" | "paused" | "hidden";
 /** Background → content: complete, authoritative control-overlay state. */
 export const OVERLAY_AGENT_STATE = "bh-agent-overlay-state";
 
-export interface OverlayAgentStateMessage {
+/**
+ * Orders the overlay state and reset messages one background worker sends:
+ * `generation` counts them, and `epoch` names the worker instance, whose
+ * counter starts over when the worker restarts.
+ */
+export interface OverlayVersion {
+  epoch: string;
+  generation: number;
+}
+
+export interface OverlayAgentStateMessage extends OverlayVersion {
   type: typeof OVERLAY_AGENT_STATE;
   sessionId: string | null;
   mode: OverlayMode;
-  generation: number;
 }
 
 export interface OverlayInterruptRequest {
@@ -95,17 +104,58 @@ export interface OverlayAutomationBypassMessage {
 /** Background → content: clear overlays that belong only inside an Agent tab. */
 export const OVERLAY_AGENT_OVERLAY_RESET = "bh-agent-overlay-reset";
 
-export interface OverlayAgentOverlayResetMessage {
+export interface OverlayAgentOverlayResetMessage extends OverlayVersion {
   type: typeof OVERLAY_AGENT_OVERLAY_RESET;
   sessionId: string;
 }
 
+function isOverlayVersion(candidate: { epoch?: unknown; generation?: unknown }): boolean {
+  return typeof candidate.epoch === "string" && typeof candidate.generation === "number";
+}
+
 export function isOverlayAgentOverlayResetMessage(
   message: unknown,
 ): message is OverlayAgentOverlayResetMessage {
   if (!message || typeof message !== "object") return false;
-  const candidate = message as { type?: unknown; sessionId?: unknown };
-  return candidate.type === OVERLAY_AGENT_OVERLAY_RESET && typeof candidate.sessionId === "string";
+  const candidate = message as {
+    type?: unknown;
+    sessionId?: unknown;
+    epoch?: unknown;
+    generation?: unknown;
+  };
+  return (
+    candidate.type === OVERLAY_AGENT_OVERLAY_RESET &&
+    typeof candidate.sessionId === "string" &&
+    isOverlayVersion(candidate)
+  );
+}
+
+/**
+ * Page side: admits overlay state and reset messages in the order the
+ * background worker sent them, dropping one that arrives after a newer one.
+ * A late `control` state after a newer `hidden` state or reset would
+ * otherwise pin the full-viewport blocker on the page. A message from
+ * another worker instance is newer, as a worker never outlives its
+ * successor, and messages from the instance it replaced stay dropped.
+ */
+export class OverlayVersionGate {
+  private last: OverlayVersion | null = null;
+  private readonly retiredEpochs = new Set<string>();
+
+  /** Whether the message with this version is current; records it if so. */
+  admit(version: OverlayVersion): boolean {
+    const last = this.last;
+    if (last && version.epoch === last.epoch) {
+      // The same message again, such as a re-sent state, is still current.
+      if (version.generation < last.generation) return false;
+    } else if (this.retiredEpochs.has(version.epoch)) {
+      return false;
+    } else if (last) {
+      this.retiredEpochs.add(last.epoch);
+    }
+    this.last = { epoch: version.epoch, generation: version.generation };
+    return true;
+  }
 }
 
 export function isOverlayAgentStateMessage(message: unknown): message is OverlayAgentStateMessage {
@@ -114,6 +164,7 @@ export function isOverlayAgentStateMessage(message: unknown): message is Overlay
     type?: unknown;
     sessionId?: unknown;
     mode?: unknown;
+    epoch?: unknown;
     generation?: unknown;
   };
   return (
@@ -123,7 +174,7 @@ export function isOverlayAgentStateMessage(message: unknown): message is Overlay
       candidate.mode === "interrupting" ||
       candidate.mode === "paused" ||
       candidate.mode === "hidden") &&
-    typeof candidate.generation === "number"
+    isOverlayVersion(candidate)
   );
 }
 

apps/extension/src/lib/overlay-version.ts

@@ -0,0 +1,16 @@
+import type { OverlayVersion } from "./overlay-bridge";
+
+// Names this worker instance. A restarted worker counts from zero again, and
+// its new epoch tells a page that kept the old counter that it is not stale.
+const epoch = crypto.randomUUID();
+let generation = 0;
+
+/**
+ * The version of the next overlay state or reset message the background
+ * worker sends to a page. Every message gets its own, so the page can order
+ * any two of them, including two different states for the same change.
+ */
+export function nextOverlayVersion(): OverlayVersion {
+  generation += 1;
+  return { epoch, generation };
+}

apps/extension/src/session-manager/__tests__/task-popups.test.ts

@@ -194,6 +194,8 @@ it.each([
   expect(sendMessage).toHaveBeenCalledWith(20, {
     type: "bh-agent-overlay-reset",
     sessionId: "one",
+    epoch: expect.any(String),
+    generation: expect.any(Number),
   });
 });
 

apps/extension/src/tools/tabs.ts

@@ -8,6 +8,7 @@ import {
   OVERLAY_AGENT_OVERLAY_RESET,
   type OverlayAgentOverlayResetMessage,
 } from "@/lib/overlay-bridge";
+import { nextOverlayVersion } from "@/lib/overlay-version";
 import {
   isAgentControlledTab,
   type SessionContext,
@@ -182,6 +183,7 @@ export const chromeAgentOverlayResetApi: AgentOverlayResetApi = {
     const message: OverlayAgentOverlayResetMessage = {
       type: OVERLAY_AGENT_OVERLAY_RESET,
       sessionId,
+      ...nextOverlayVersion(),
     };
     await chrome.tabs.sendMessage(tabId, message);
   },