Fix Flow key loss after MV3 service worker restart

extension/background.js

@@ -62,17 +62,35 @@ function broadcastRequestLog() {
 
 // ─── Startup ────────────────────────────────────────────────
 
-chrome.runtime.onInstalled.addListener(init);
-chrome.runtime.onStartup.addListener(init);
+let initializationPromise = null;
+
+chrome.runtime.onInstalled.addListener(() => {
+  void ensureInitialized();
+});
+chrome.runtime.onStartup.addListener(() => {
+  void ensureInitialized();
+});
 chrome.alarms.onAlarm.addListener(async (alarm) => {
+  await ensureInitialized();
   if (alarm.name === 'reconnect') connectToAgent();
   if (alarm.name === 'keepAlive') keepAlive();
   if (alarm.name === 'token-refresh') {
     await captureTokenFromFlowTab();
   }
 });
 
-async function init() {
+function ensureInitialized() {
+  if (!initializationPromise) {
+    initializationPromise = initialize().catch((error) => {
+      initializationPromise = null;
+      console.error('[FlowAgent] Initialization failed', error);
+      throw error;
+    });
+  }
+  return initializationPromise;
+}
+
+async function initialize() {
   const data = await chrome.storage.local.get(['flowKey', 'metrics', 'callbackSecret']);
   if (data.flowKey) flowKey = data.flowKey;
   if (data.metrics) Object.assign(metrics, data.metrics);
@@ -81,6 +99,10 @@ async function init() {
   chrome.alarms.create('keepAlive', { periodInMinutes: 0.4 });
 }
 
+// MV3 workers can be suspended and restarted without onStartup firing.
+// Rehydrate the persisted Flow key on every worker start.
+void ensureInitialized();
+
 // ─── Token Capture ──────────────────────────────────────────
 
 chrome.webRequest.onBeforeSendHeaders.addListener(

tests/extension_mv3_bootstrap.test.cjs

@@ -0,0 +1,112 @@
+const assert = require('node:assert/strict');
+const fs = require('node:fs');
+const path = require('node:path');
+const vm = require('node:vm');
+
+const source = fs.readFileSync(
+  path.join(__dirname, '..', 'extension', 'background.js'),
+  'utf8',
+);
+
+const lifecycleListeners = { alarm: [], installed: [], startup: [] };
+const sockets = [];
+let storageReads = 0;
+
+class FakeWebSocket {
+  static CONNECTING = 0;
+  static OPEN = 1;
+
+  constructor(url) {
+    this.url = url;
+    this.readyState = FakeWebSocket.CONNECTING;
+    this.messages = [];
+    sockets.push(this);
+  }
+
+  send(message) {
+    this.messages.push(JSON.parse(message));
+  }
+}
+
+function event(bucket) {
+  return {
+    addListener(listener) {
+      if (bucket) bucket.push(listener);
+    },
+  };
+}
+
+const chrome = {
+  action: { setBadgeBackgroundColor() {}, setBadgeText() {} },
+  alarms: { clear() {}, create() {}, onAlarm: event(lifecycleListeners.alarm) },
+  runtime: {
+    onInstalled: event(lifecycleListeners.installed),
+    onMessage: event(),
+    onStartup: event(lifecycleListeners.startup),
+    sendMessage: async () => {},
+  },
+  scripting: { executeScript: async () => {} },
+  storage: {
+    local: {
+      async get() {
+        storageReads += 1;
+        return {
+          callbackSecret: 'persisted-secret',
+          flowKey: 'persisted-flow-key',
+          metrics: { tokenCapturedAt: 1234 },
+        };
+      },
+      async set() {},
+    },
+  },
+  tabs: {
+    create: async () => ({}),
+    query: async () => [],
+    sendMessage: async () => {},
+    update: async () => {},
+  },
+  webRequest: { onBeforeSendHeaders: event() },
+};
+
+const context = vm.createContext({
+  URL,
+  WebSocket: FakeWebSocket,
+  chrome,
+  clearInterval() {},
+  clearTimeout() {},
+  console,
+  fetch: async () => ({ ok: true }),
+  navigator: { userAgent: 'FlowkitBootstrapTest/1.0' },
+  setInterval() { return 1; },
+  setTimeout() { return 1; },
+});
+
+vm.runInContext(source, context, { filename: 'background.js' });
+
+setImmediate(async () => {
+  await Promise.resolve();
+  await Promise.resolve();
+
+  assert.equal(storageReads, 1, 'cold worker start must hydrate storage exactly once');
+  assert.equal(sockets.length, 1, 'cold worker start must connect after hydration');
+
+  await lifecycleListeners.startup[0]();
+  await lifecycleListeners.installed[0]();
+  await lifecycleListeners.alarm[0]({ name: 'keepAlive' });
+  assert.equal(storageReads, 1, 'lifecycle events must reuse initialization');
+  assert.equal(sockets.length, 1, 'lifecycle events must not duplicate the socket');
+
+  const socket = sockets[0];
+  socket.readyState = FakeWebSocket.OPEN;
+  socket.onopen();
+
+  assert.equal(socket.messages[0].type, 'extension_ready');
+  assert.equal(socket.messages[0].flowKeyPresent, true);
+  assert.ok(socket.messages[0].tokenAge > 0);
+  assert.deepEqual(socket.messages[1], {
+    type: 'token_captured',
+    flowKey: 'persisted-flow-key',
+  });
+
+  console.log('Flowkit MV3 cold-start bootstrap regression test passed');
+});