diff --git a/extension/background.js b/extension/background.js index fb1b1b415..ab5b00156 100644 --- a/extension/background.js +++ b/extension/background.js @@ -62,9 +62,16 @@ 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') { @@ -72,7 +79,18 @@ chrome.alarms.onAlarm.addListener(async (alarm) => { } }); -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( diff --git a/tests/extension_mv3_bootstrap.test.cjs b/tests/extension_mv3_bootstrap.test.cjs new file mode 100644 index 000000000..86a12f54d --- /dev/null +++ b/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'); +});