fix(browser): clear relay deadline on socket close (#123343)

extensions/browser/chrome-extension/background.access-mode.test.ts

@@ -1,5 +1,6 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import {
+  cleanupBackgroundHarnesses,
   loadBackground,
   TEST_RELAY_KEY,
   REPLACEMENT_TEST_RELAY_KEY,
@@ -13,7 +14,8 @@ describe("relay command authorization", () => {
     vi.resetModules();
   });
 
-  afterEach(() => {
+  afterEach(async () => {
+    await cleanupBackgroundHarnesses();
     vi.unstubAllGlobals();
   });
 

extensions/browser/chrome-extension/background.js

@@ -87,6 +87,7 @@ async function requireAutomationAllowed() {
 }
 
 function closeRelaySocket() {
+  clearRelayOpeningDeadline();
   const socket = relayWs;
   if (!socket) {
     return;
@@ -113,7 +114,6 @@ async function reconcilePairingInvalidation() {
     return;
   }
   reconciledPairingInvalidationRevision = pairingConfigStore.invalidationRevision;
-  clearRelayOpeningDeadline();
   await syncTabsToRelay();
   closeRelaySocket();
   setBadge("off");
@@ -440,7 +440,6 @@ async function connectRelay(isConnectionAllowed = () => true) {
   await tabAccessReady;
   if (retiredCopilotCustodyBlocked) {
     tabAccessPolicy.setEnabled(false);
-    clearRelayOpeningDeadline();
     closeRelaySocket();
     setBadge("off");
     return;
@@ -521,13 +520,6 @@ async function connectRelay(isConnectionAllowed = () => true) {
 }
 
 function handleRelayOpeningDeadline() {
-  // Unit-test module isolation can outlive the mocked Chrome global. The real
-  // MV3 worker always has chrome; a detached test timer has no owner to mutate.
-  if (typeof chrome === "undefined") {
-    relayOpeningDeadlineAt = 0;
-    relayOpeningDeadlineTimer = null;
-    return;
-  }
   const ws = relayWs;
   if (!ws) {
     clearRelayOpeningDeadline();
@@ -621,7 +613,6 @@ const handlePopupMessage = createPopupMessageHandler({
   runAccessMutation,
   detachAllDebuggerSessions,
   syncTabsToRelay,
-  clearRelayOpeningDeadline,
   closeRelaySocket,
   connectRelay,
   setBadge,

extensions/browser/chrome-extension/background.test-harness.ts

@@ -14,6 +14,11 @@ export const TEST_RELAY_KEY = relayTestKey(1);
 export const REPLACEMENT_TEST_RELAY_KEY = relayTestKey(2);
 const PAIRING_CONFIG_KEYS = ["relayUrl", "token", "pairingStatus"];
 const RETIRED_CUSTODY_BLOCKED_KEY = "retiredCopilotCustodyBlockedV1";
+const backgroundCleanups = new Set<() => Promise<void>>();
+
+export async function cleanupBackgroundHarnesses(): Promise<void> {
+  await Promise.all([...backgroundCleanups].map(async (cleanup) => await cleanup()));
+}
 
 export type RetiredStorageFailureStage =
   | "marker_set"
@@ -416,6 +421,14 @@ export async function loadBackground({
   ) {
     throw new Error("expected background worker lifecycle listeners");
   }
+  const lifecycleMessageListener = messageListener;
+  const cleanup = async () => {
+    backgroundCleanups.delete(cleanup);
+    await new Promise<void>((resolve) => {
+      lifecycleMessageListener({ type: "unpair" }, {}, () => resolve());
+    });
+  };
+  backgroundCleanups.add(cleanup);
   return {
     alarmListener,
     clearAlarm,

extensions/browser/chrome-extension/background.test.ts

@@ -1,5 +1,6 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import {
+  cleanupBackgroundHarnesses,
   loadBackground,
   TEST_RELAY_KEY,
   REPLACEMENT_TEST_RELAY_KEY,
@@ -22,7 +23,8 @@ describe("native extension bootstrap", () => {
     vi.resetModules();
   });
 
-  afterEach(() => {
+  afterEach(async () => {
+    await cleanupBackgroundHarnesses();
     vi.unstubAllGlobals();
   });
 
@@ -347,7 +349,8 @@ describe("relay pairing and authentication", () => {
     vi.resetModules();
   });
 
-  afterEach(() => {
+  afterEach(async () => {
+    await cleanupBackgroundHarnesses();
     vi.unstubAllGlobals();
   });
 
@@ -366,6 +369,23 @@ describe("relay pairing and authentication", () => {
     expect(JSON.stringify(harness.relaySockets[0]?.protocols)).not.toContain(TEST_RELAY_KEY);
   });
 
+  it("cancels the opening deadline when the socket-close fallback ends the relay", async () => {
+    const harness = await loadBackground({ relayNegotiatedProtocol: "unsupported" });
+    const socket = harness.relaySockets[0];
+    if (!socket) {
+      throw new Error("expected relay socket");
+    }
+    harness.clearAlarm.mockClear();
+    socket.close.mockImplementationOnce(() => {
+      throw new Error("socket close failed");
+    });
+
+    socket.open();
+
+    await vi.waitFor(() => expect(socket.close).toHaveBeenCalledTimes(2));
+    expect(harness.clearAlarm).toHaveBeenCalledOnce();
+  });
+
   it("revokes synchronously while an older manual pair is stalled", async () => {
     const harness = await loadBackground({
       initialTabs: [{ id: 131, url: "https://example.com/paired", groupId: 7 }],

extensions/browser/chrome-extension/modules/popup-background.js

@@ -38,7 +38,6 @@ export function createPopupMessageHandler({
   runAccessMutation,
   detachAllDebuggerSessions,
   syncTabsToRelay,
-  clearRelayOpeningDeadline,
   closeRelaySocket,
   connectRelay,
   setBadge,
@@ -71,7 +70,6 @@ export function createPopupMessageHandler({
     }
     const generation = ++pairingGeneration;
     suspendRelayConnections();
-    clearRelayOpeningDeadline();
     closeRelaySocket();
     await accessReady;
     assertPairingCurrent(generation);
@@ -81,7 +79,6 @@ export function createPopupMessageHandler({
         return;
       }
       suspendRelayConnections();
-      clearRelayOpeningDeadline();
       closeRelaySocket();
       const normalizedMode =
         accessMode === ACCESS_MODE_SELECTED ? ACCESS_MODE_SELECTED : ACCESS_MODE_ALL;
@@ -107,7 +104,6 @@ export function createPopupMessageHandler({
       resumeRelayConnections();
       await connectRelay(() => generation === pairingGeneration);
       if (generation !== pairingGeneration) {
-        clearRelayOpeningDeadline();
         closeRelaySocket();
         setBadge("off");
         assertPairingCurrent(generation);
@@ -123,13 +119,11 @@ export function createPopupMessageHandler({
     policy.invalidateAll();
     suspendRelayConnections();
     resetRelayState();
-    clearRelayOpeningDeadline();
     closeRelaySocket();
     setBadge("off");
     await accessReady;
     policy.setEnabled(false);
     policy.invalidateAll();
-    clearRelayOpeningDeadline();
     closeRelaySocket();
     setBadge("off");
     await runAccessMutation(async () => {
@@ -142,7 +136,6 @@ export function createPopupMessageHandler({
       await detaching;
       await discardRetiredCopilotCustody();
       resetRelayState();
-      clearRelayOpeningDeadline();
       closeRelaySocket();
       setBadge("off");
     });