Fix shared worker resumption after tab suspend (#7656)

* [AI] Fix SharedWorker tab resume recovery

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* [AI] Fix SharedWorker reload readiness

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Add release notes

* Update packages/desktop-client/src/shared-browser-server-core.ts

Co-authored-by: Matiss Janis Aboltins <matiss@mja.lv>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Matiss Janis Aboltins <matiss@mja.lv>

packages/desktop-client/src/browser-preload.js

@@ -27,6 +27,13 @@ let worker = null;
 // The regular Worker running the backend, created only on the leader tab
 let localBackendWorker = null;
 
+function terminateLocalBackendWorker() {
+  if (localBackendWorker) {
+    localBackendWorker.terminate();
+    localBackendWorker = null;
+  }
+}
+
 /**
  * WorkerBridge wraps a SharedWorker port and presents a Worker-like interface
  * (onmessage, postMessage, addEventListener, start) to the connection layer.
@@ -43,9 +50,22 @@ class WorkerBridge {
     this._onmessage = null;
     this._listeners = [];
     this._started = false;
+    this._isInitialized = false;
+    this._currentBudgetId = null;
+    this._wasHidden = document.visibilityState === 'hidden';
+
+    this._onVisibilityChange = () => {
+      if (document.visibilityState === 'hidden') {
+        this._wasHidden = true;
+      } else if (this._wasHidden) {
+        this._wasHidden = false;
+        this._resumeAssociation();
+      }
+    };
 
     // Listen for all messages from the SharedWorker port
     sharedPort.addEventListener('message', e => this._onSharedMessage(e));
+    document.addEventListener('visibilitychange', this._onVisibilityChange);
   }
 
   set onmessage(handler) {
@@ -109,10 +129,7 @@ class WorkerBridge {
     // show-budgets normally.
     if (msg && msg.type === '__close-and-transfer') {
       console.log('[WorkerBridge] Leadership transferred — terminating Worker');
-      if (localBackendWorker) {
-        localBackendWorker.terminate();
-        localBackendWorker = null;
-      }
+      this._applyRole('UNASSIGNED', null);
       // Only dispatch a synthetic reply if there's an actual close-budget
       // request to complete. When requestId is null the eviction was
       // triggered externally (e.g. another tab deleted this budget).
@@ -126,6 +143,7 @@ class WorkerBridge {
 
     // Role change notification
     if (msg && msg.type === '__role-change') {
+      this._applyRole(msg.role, msg.budgetId ?? null);
       console.log(
         `[WorkerBridge] Role: ${msg.role}${msg.budgetId ? ` (budget: ${msg.budgetId})` : ''}`,
       );
@@ -146,13 +164,47 @@ class WorkerBridge {
     }
 
     // Everything else goes to the connection layer
+    if (msg && msg.type === 'push' && msg.name === 'show-budgets') {
+      this._applyRole('UNASSIGNED', null);
+    }
     this._dispatch(event);
   }
 
-  _createLocalWorker(initMsg, budgetToRestore, pendingMsg) {
-    if (localBackendWorker) {
-      localBackendWorker.terminate();
+  markInitialized() {
+    this._isInitialized = true;
+  }
+
+  _normalizeBudgetId(budgetId) {
+    if (
+      typeof budgetId === 'string' &&
+      budgetId.length > 0 &&
+      !budgetId.startsWith('__')
+    ) {
+      return budgetId;
     }
+    return null;
+  }
+
+  _applyRole(role, budgetId) {
+    this._currentBudgetId = this._normalizeBudgetId(budgetId);
+
+    if (role !== 'LEADER') {
+      terminateLocalBackendWorker();
+    }
+  }
+
+  _resumeAssociation() {
+    if (!this._isInitialized) {
+      return;
+    }
+    this._sharedPort.postMessage({
+      type: '__resume-tab',
+      budgetId: this._currentBudgetId,
+    });
+  }
+
+  _createLocalWorker(initMsg, budgetToRestore, pendingMsg) {
+    terminateLocalBackendWorker();
     localBackendWorker = new Worker(backendWorkerUrl);
     initSQLBackend(localBackendWorker);
 
@@ -238,10 +290,12 @@ function createBackendWorker() {
           'SharedArrayBufferOverride',
         ),
       });
+      worker.markInitialized();
 
-      window.addEventListener('beforeunload', () => {
+      const notifyTabClosing = () => {
         sharedPort.postMessage({ type: 'tab-closing' });
-      });
+      };
+      window.addEventListener('beforeunload', notifyTabClosing);
 
       return;
     } catch (e) {

packages/desktop-client/src/shared-browser-server-core.test.ts

@@ -380,7 +380,7 @@ describe('SharedWorker coordinator', () => {
 
   describe('tab disconnection', () => {
     it('leader disconnect promotes follower', () => {
-      setupBudgetGroup(coordinator, 'budget-1');
+      const leader = setupBudgetGroup(coordinator, 'budget-1');
 
       const follower = connectTab(coordinator);
       sendInit(follower);
@@ -391,10 +391,6 @@ describe('SharedWorker coordinator', () => {
       });
       follower.postMessage.mockClear();
 
-      // Find current leader to disconnect it
-      const group = coordinator.getState().budgetGroups.get('budget-1');
-      const leader = group.leaderPort as MockPort;
-
       // Leader closes tab
       sendMsg(leader, { type: 'tab-closing' });
 
@@ -481,6 +477,116 @@ describe('SharedWorker coordinator', () => {
     });
   });
 
+  describe('__resume-tab', () => {
+    it('keeps the lobby leader attached during startup resume signals', () => {
+      const leader = connectTab(coordinator);
+      sendInit(leader);
+      leader.postMessage.mockClear();
+
+      sendMsg(leader, { type: '__resume-tab', budgetId: null });
+
+      expect(coordinator.getState().portToBudget.get(leader)).toBe('__lobby');
+      expect(coordinator.getState().unassignedPorts.has(leader)).toBe(false);
+      expect(
+        coordinator.getState().budgetGroups.get('__lobby').leaderPort,
+      ).toBe(leader);
+      expect(leader.postMessage).not.toHaveBeenCalledWith(
+        expect.objectContaining({
+          type: '__role-change',
+          role: 'UNASSIGNED',
+        }),
+      );
+    });
+
+    it('does not route orphaned ports through another live budget before they resume', () => {
+      const orphanedLeader = setupBudgetGroup(coordinator, 'budget-1');
+      const liveLeader = setupBudgetGroup(coordinator, 'budget-2');
+      orphanedLeader.postMessage.mockClear();
+      liveLeader.postMessage.mockClear();
+
+      vi.advanceTimersByTime(10_000);
+      sendMsg(liveLeader, { type: '__heartbeat-pong' });
+      vi.advanceTimersByTime(10_000);
+
+      sendMsg(orphanedLeader, { id: 'req-orphan', name: 'get-budgets' });
+
+      expect(liveLeader.postMessage).not.toHaveBeenCalledWith(
+        expect.objectContaining({
+          type: '__to-worker',
+          msg: expect.objectContaining({ name: 'get-budgets' }),
+        }),
+      );
+    });
+
+    it('re-elects an orphaned solo tab as leader and restores its budget', () => {
+      const leader = setupBudgetGroup(coordinator, 'budget-1');
+      leader.postMessage.mockClear();
+
+      vi.advanceTimersByTime(10_000);
+      vi.advanceTimersByTime(10_000);
+
+      sendMsg(leader, { type: '__resume-tab', budgetId: 'budget-1' });
+
+      expect(coordinator.getState().connectedPorts.includes(leader)).toBe(true);
+      expect(coordinator.getState().portToBudget.get(leader)).toBe('budget-1');
+      expect(
+        coordinator.getState().budgetGroups.get('budget-1').leaderPort,
+      ).toBe(leader);
+      expect(leader.postMessage).toHaveBeenCalledWith(
+        expect.objectContaining({
+          type: '__role-change',
+          role: 'LEADER',
+          budgetId: 'budget-1',
+        }),
+      );
+      expect(leader.postMessage).toHaveBeenCalledWith(
+        expect.objectContaining({
+          type: '__become-leader',
+          budgetToRestore: 'budget-1',
+        }),
+      );
+    });
+
+    it('reattaches an orphaned tab to an existing budget group as follower', () => {
+      const leader = setupBudgetGroup(coordinator, 'budget-1');
+
+      const follower = connectTab(coordinator);
+      sendInit(follower);
+      sendMsg(follower, {
+        id: 'lb-f',
+        name: 'load-budget',
+        args: { id: 'budget-1' },
+      });
+      follower.postMessage.mockClear();
+
+      vi.advanceTimersByTime(10_000);
+      sendMsg(leader, { type: '__heartbeat-pong' });
+      vi.advanceTimersByTime(10_000);
+
+      sendMsg(follower, { type: '__resume-tab', budgetId: 'budget-1' });
+
+      expect(coordinator.getState().connectedPorts.includes(follower)).toBe(
+        true,
+      );
+      expect(
+        coordinator
+          .getState()
+          .budgetGroups.get('budget-1')
+          .followers.has(follower),
+      ).toBe(true);
+      expect(follower.postMessage).toHaveBeenCalledWith(
+        expect.objectContaining({
+          type: '__role-change',
+          role: 'FOLLOWER',
+          budgetId: 'budget-1',
+        }),
+      );
+      expect(follower.postMessage).toHaveBeenCalledWith(
+        expect.objectContaining({ type: 'connect' }),
+      );
+    });
+  });
+
   // ── Worker message routing ──────────────────────────────────────────
 
   describe('Worker message routing', () => {
@@ -789,6 +895,98 @@ describe('SharedWorker coordinator', () => {
       expect(group.requestNames.get('restore-1')).toBe('load-budget');
       expect(group.requestBudgetIds.get('restore-1')).toBe('budget-1');
     });
+
+    it('waits to broadcast connect until a promoted leader finishes restoring', () => {
+      const leader = setupBudgetGroup(coordinator, 'budget-1');
+
+      const follower = connectTab(coordinator);
+      sendInit(follower);
+      sendMsg(follower, {
+        id: 'lb-f',
+        name: 'load-budget',
+        args: { id: 'budget-1' },
+      });
+      follower.postMessage.mockClear();
+
+      sendMsg(leader, { id: 'cb-leader', name: 'close-budget' });
+      sendMsg(follower, {
+        type: '__track-restore',
+        requestId: '__restore-budget',
+        budgetId: 'budget-1',
+      });
+
+      const reloaded = connectTab(coordinator);
+      sendInit(reloaded);
+      reloaded.postMessage.mockClear();
+
+      sendMsg(follower, {
+        type: '__from-worker',
+        msg: { type: 'connect' },
+      });
+
+      expect(
+        coordinator.getState().budgetGroups.get('budget-1').backendConnected,
+      ).toBe(false);
+      expect(reloaded.postMessage).not.toHaveBeenCalledWith(
+        expect.objectContaining({ type: 'connect' }),
+      );
+
+      sendMsg(follower, {
+        type: '__from-worker',
+        msg: { type: 'reply', id: '__restore-budget', result: {} },
+      });
+
+      expect(
+        coordinator.getState().budgetGroups.get('budget-1').backendConnected,
+      ).toBe(true);
+      expect(reloaded.postMessage).toHaveBeenCalledWith(
+        expect.objectContaining({ type: 'connect' }),
+      );
+    });
+
+    it('still broadcasts connect if restore finishes before the worker connect event', () => {
+      const leader = setupBudgetGroup(coordinator, 'budget-1');
+
+      const follower = connectTab(coordinator);
+      sendInit(follower);
+      sendMsg(follower, {
+        id: 'lb-f',
+        name: 'load-budget',
+        args: { id: 'budget-1' },
+      });
+
+      sendMsg(leader, { id: 'cb-leader', name: 'close-budget' });
+      sendMsg(follower, {
+        type: '__track-restore',
+        requestId: '__restore-budget',
+        budgetId: 'budget-1',
+      });
+
+      const reloaded = connectTab(coordinator);
+      sendInit(reloaded);
+      reloaded.postMessage.mockClear();
+
+      sendMsg(follower, {
+        type: '__from-worker',
+        msg: { type: 'reply', id: '__restore-budget', result: {} },
+      });
+
+      expect(reloaded.postMessage).not.toHaveBeenCalledWith(
+        expect.objectContaining({ type: 'connect' }),
+      );
+
+      sendMsg(follower, {
+        type: '__from-worker',
+        msg: { type: 'connect' },
+      });
+
+      expect(
+        coordinator.getState().budgetGroups.get('budget-1').backendConnected,
+      ).toBe(true);
+      expect(reloaded.postMessage).toHaveBeenCalledWith(
+        expect.objectContaining({ type: 'connect' }),
+      );
+    });
   });
 
   // ── Multiple budgets ────────────────────────────────────────────────

packages/desktop-client/src/shared-browser-server-core.ts

@@ -19,6 +19,8 @@ type BudgetGroup = {
   leaderPort: CoordinatorPort;
   followers: Set<CoordinatorPort>;
   backendConnected: boolean;
+  pendingConnect: boolean;
+  restoreBudgetId: string | null;
   requestToPort: Map<string, CoordinatorPort>;
   requestNames: Map<string, string>;
   requestBudgetIds: Map<string, string>;
@@ -89,12 +91,22 @@ export function createCoordinator({
       leaderPort,
       followers: new Set(),
       backendConnected: false,
+      pendingConnect: false,
+      restoreBudgetId: null,
       requestToPort: new Map(),
       requestNames: new Map(),
       requestBudgetIds: new Map(),
     };
   }
 
+  function broadcastConnect(budgetId: string) {
+    const connectMsg = { type: 'connect' };
+    broadcastToAllInGroup(budgetId, connectMsg);
+    for (const port of unassignedPorts) {
+      port.postMessage(connectMsg);
+    }
+  }
+
   function logState(action: string) {
     const groups: string[] = [];
     for (const [bid, g] of budgetGroups) {
@@ -105,6 +117,110 @@ export function createCoordinator({
     );
   }
 
+  function isTrackedPort(port: CoordinatorPort) {
+    return connectedPorts.includes(port);
+  }
+
+  function ensureTrackedPort(port: CoordinatorPort) {
+    if (!isTrackedPort(port)) {
+      connectedPorts.push(port);
+    }
+  }
+
+  function hasConnectedBackend() {
+    for (const [, group] of budgetGroups) {
+      if (group.backendConnected) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  function isGroupMember(group: BudgetGroup, port: CoordinatorPort) {
+    return group.leaderPort === port || group.followers.has(port);
+  }
+
+  function movePortToUnassigned(port: CoordinatorPort) {
+    const prevBudget = portToBudget.get(port);
+    if (prevBudget) {
+      removePortFromGroup(port, prevBudget);
+    }
+
+    portToBudget.delete(port);
+    unassignedPorts.add(port);
+  }
+
+  function restoreUnassignedPort(port: CoordinatorPort) {
+    movePortToUnassigned(port);
+    port.postMessage({ type: '__role-change', role: 'UNASSIGNED' });
+
+    if (hasConnectedBackend()) {
+      port.postMessage({ type: 'connect' });
+    }
+  }
+
+  function resumePort(port: CoordinatorPort, budgetId?: string | null) {
+    const normalizedBudgetId = budgetId || null;
+    const wasTracked = isTrackedPort(port);
+    const currentBudget = portToBudget.get(port);
+    const currentGroup = currentBudget ? budgetGroups.get(currentBudget) : null;
+    const alreadyAttached =
+      !!normalizedBudgetId &&
+      currentBudget === normalizedBudgetId &&
+      !!currentGroup &&
+      isGroupMember(currentGroup, port);
+    const alreadyUnassigned =
+      !normalizedBudgetId &&
+      !currentBudget &&
+      wasTracked &&
+      unassignedPorts.has(port);
+    const alreadyOnTemporaryGroup =
+      !normalizedBudgetId &&
+      !!currentBudget &&
+      !!currentGroup &&
+      currentBudget.startsWith('__') &&
+      isGroupMember(currentGroup, port);
+
+    ensureTrackedPort(port);
+    pendingPongs.delete(port);
+
+    if (alreadyAttached) {
+      logState(`Tab resumed on budget "${normalizedBudgetId}"`);
+      return;
+    }
+
+    if (alreadyUnassigned) {
+      logState('Tab resume confirmed while unassigned');
+      return;
+    }
+
+    if (alreadyOnTemporaryGroup) {
+      logState(`Tab resumed on coordinator group "${currentBudget}"`);
+      return;
+    }
+
+    if (!normalizedBudgetId) {
+      if (budgetGroups.size === 0) {
+        electLeader('__lobby', port);
+        logState('Tab resumed into lobby');
+      } else {
+        restoreUnassignedPort(port);
+        logState('Tab resumed as unassigned');
+      }
+      return;
+    }
+
+    const existingGroup = budgetGroups.get(normalizedBudgetId);
+    if (existingGroup) {
+      addFollower(normalizedBudgetId, port);
+      logState(`Tab resumed on budget "${normalizedBudgetId}" as follower`);
+      return;
+    }
+
+    electLeader(normalizedBudgetId, port, normalizedBudgetId);
+    logState(`Tab resumed on budget "${normalizedBudgetId}" as leader`);
+  }
+
   function broadcastToGroup(
     budgetId: string,
     msg: unknown,
@@ -194,10 +310,15 @@ export function createCoordinator({
     } else {
       group.leaderPort = port;
       group.backendConnected = false;
+      group.pendingConnect = false;
+      group.restoreBudgetId = budgetToRestore || null;
       group.requestToPort.clear();
       group.requestNames.clear();
       group.requestBudgetIds.clear();
     }
+    if (!group.restoreBudgetId && budgetToRestore) {
+      group.restoreBudgetId = budgetToRestore;
+    }
     const prevBudget = portToBudget.get(port);
     if (prevBudget && prevBudget !== budgetId) {
       removePortFromGroup(port, prevBudget);
@@ -323,6 +444,15 @@ export function createCoordinator({
       );
     }
 
+    if (oldGroup.restoreBudgetId === newBudgetId) {
+      oldGroup.restoreBudgetId = null;
+      if (oldGroup.pendingConnect) {
+        oldGroup.backendConnected = true;
+        oldGroup.pendingConnect = false;
+        broadcastConnect(newBudgetId);
+      }
+    }
+
     logState(`Budget "${newBudgetId}" ready`);
   }
 
@@ -401,6 +531,14 @@ export function createCoordinator({
           return;
         }
 
+        if (msg.type === '__resume-tab') {
+          resumePort(
+            port,
+            typeof msg.budgetId === 'string' ? msg.budgetId : null,
+          );
+          return;
+        }
+
         if (msg.type === '__heartbeat-pong') {
           pendingPongs.delete(port);
           return;
@@ -413,14 +551,7 @@ export function createCoordinator({
           if (lastAppInitFailure) {
             port.postMessage(lastAppInitFailure);
           } else {
-            let anyConnected = false;
-            for (const [, g] of budgetGroups) {
-              if (g.backendConnected) {
-                anyConnected = true;
-                break;
-              }
-            }
-            if (anyConnected) {
+            if (hasConnectedBackend()) {
               port.postMessage({ type: '__role-change', role: 'UNASSIGNED' });
               port.postMessage({ type: 'connect' });
             } else if (budgetGroups.size > 0) {
@@ -475,10 +606,11 @@ export function createCoordinator({
               group.requestNames.delete(workerMsg.id as string);
             }
           } else if (workerMsg.type === 'connect') {
-            group.backendConnected = true;
-            broadcastToAllInGroup(portBudget!, workerMsg);
-            for (const p of unassignedPorts) {
-              p.postMessage(workerMsg);
+            if (group.restoreBudgetId) {
+              group.pendingConnect = true;
+            } else {
+              group.backendConnected = true;
+              broadcastConnect(portBudget!);
             }
           } else if (workerMsg.type === 'app-init-failure') {
             lastAppInitFailure = workerMsg;
@@ -493,6 +625,7 @@ export function createCoordinator({
 
         if (msg.type === '__track-restore') {
           if (group) {
+            group.restoreBudgetId = msg.budgetId as string;
             group.requestToPort.set(msg.requestId as string, port);
             group.requestNames.set(msg.requestId as string, 'load-budget');
             group.requestBudgetIds.set(
@@ -689,7 +822,7 @@ export function createCoordinator({
         // ── Default: track and forward to leader ───────────────────
 
         let targetGroup = group;
-        if (!targetGroup) {
+        if (!targetGroup && unassignedPorts.has(port) && isTrackedPort(port)) {
           for (const [, g] of budgetGroups) {
             if (g.backendConnected) {
               targetGroup = g;

upcoming-release-notes/7656.md

@@ -0,0 +1,6 @@
+---
+category: Bugfixes
+authors: [jfdoming]
+---
+
+Fix shared worker resumption after tab suspend