From 5e12cfd1e054592dda2d3ab90f9bc57f52fc705f Mon Sep 17 00:00:00 2001 From: Steve Olson Date: Sat, 28 Feb 2026 12:52:00 -0500 Subject: [PATCH 1/2] fix(extension): persist and serialize per-tab runtime state Persist per-tab runtime state in chrome.storage.session so hard reloads and MV3 worker restarts do not reset This Tab toggle. Convert background state accessors to async and await them in message handlers. Serialize storage writes with a queue to prevent stale snapshot overwrites, and add regression tests for restart restore and write-order safety. --- packages/extension/src/background/messages.ts | 18 +-- .../extension/src/background/state.test.ts | 121 ++++++++++++++ packages/extension/src/background/state.ts | 147 ++++++++++++++++-- 3 files changed, 262 insertions(+), 24 deletions(-) create mode 100644 packages/extension/src/background/state.test.ts diff --git a/packages/extension/src/background/messages.ts b/packages/extension/src/background/messages.ts index 48c9861..77c7045 100644 --- a/packages/extension/src/background/messages.ts +++ b/packages/extension/src/background/messages.ts @@ -25,7 +25,7 @@ async function ensureContentScriptInjected(tabId: number): Promise { } async function reInjectContentScriptIfEnabled(tabId: number, url?: string): Promise { - const state = stateManager.getTabRuntimeState(tabId); + const state = await stateManager.getTabRuntimeState(tabId); if (!state.enabled || !isInjectableTabUrl(url)) { return; } @@ -44,7 +44,7 @@ async function reInjectContentScriptIfEnabled(tabId: number, url?: string): Prom } async function notifyTabRuntimeStateChanged(tabId: number): Promise { - const state = stateManager.getTabRuntimeState(tabId); + const state = await stateManager.getTabRuntimeState(tabId); const payload = { type: 'TAB_RUNTIME_STATE_CHANGED' as const, payload: { tabId, state }, @@ -251,7 +251,7 @@ async function handleMessage( }; } - const state = stateManager.getTabRuntimeState(tabId); + const state = await stateManager.getTabRuntimeState(tabId); console.log(`${LOG_PREFIX} ${requestId} GET_TAB_RUNTIME_STATE completed`, { durationMs: Date.now() - receivedAt, tabId, @@ -278,7 +278,7 @@ async function handleMessage( } } - const state = stateManager.setTabEnabled(tabId, message.payload.enabled); + const state = await stateManager.setTabEnabled(tabId, message.payload.enabled); await notifyTabRuntimeStateChanged(tabId); console.log(`${LOG_PREFIX} ${requestId} SET_TAB_ENABLED completed`, { @@ -297,12 +297,12 @@ async function handleMessage( } // MV3 service workers are ephemeral; recover enabled state when content is already active. - const current = stateManager.getTabRuntimeState(tabId); + const current = await stateManager.getTabRuntimeState(tabId); if (!current.enabled && message.meta?.source === 'content') { - stateManager.setTabEnabled(tabId, true); + await stateManager.setTabEnabled(tabId, true); } - const state = stateManager.setAnnotateMode(tabId, message.payload.annotateMode); + const state = await stateManager.setAnnotateMode(tabId, message.payload.annotateMode); await notifyTabRuntimeStateChanged(tabId); console.log(`${LOG_PREFIX} ${requestId} SET_ANNOTATE_MODE completed`, { @@ -321,7 +321,7 @@ async function handleMessage( console.warn(`${LOG_PREFIX} ${requestId} GET_STATE no tabId, returning default`); return { success: true, data: { isActive: false } }; } - const isActive = stateManager.getState(tabId); + const isActive = await stateManager.getState(tabId); console.log(`${LOG_PREFIX} ${requestId} GET_STATE completed`, { durationMs: Date.now() - receivedAt, tabId, @@ -338,7 +338,7 @@ async function handleMessage( console.warn(`${LOG_PREFIX} ${requestId} SET_STATE no tabId`); return { success: false, error: 'Missing tabId for SET_STATE' }; } - const state = stateManager.setAnnotateMode(tabId, message.payload.isActive); + const state = await stateManager.setAnnotateMode(tabId, message.payload.isActive); await notifyTabRuntimeStateChanged(tabId); console.log(`${LOG_PREFIX} ${requestId} SET_STATE completed`, { durationMs: Date.now() - receivedAt, diff --git a/packages/extension/src/background/state.test.ts b/packages/extension/src/background/state.test.ts new file mode 100644 index 0000000..edbd8bb --- /dev/null +++ b/packages/extension/src/background/state.test.ts @@ -0,0 +1,121 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +describe('StateManager runtime state persistence', () => { + const tabId = 42; + let storageBucket: Record; + let storageSetMock: ReturnType; + + beforeEach(() => { + vi.resetModules(); + storageBucket = {}; + storageSetMock = vi.fn(async (items: Record) => { + Object.assign(storageBucket, items); + }); + + const storageGetMock = vi.fn(async (keys?: string | string[] | Record) => { + if (typeof keys === 'string') { + return { [keys]: storageBucket[keys] }; + } + + if (Array.isArray(keys)) { + return keys.reduce>((acc, key) => { + acc[key] = storageBucket[key]; + return acc; + }, {}); + } + + return storageBucket; + }); + + Object.defineProperty(globalThis, 'chrome', { + value: { + tabs: { + onRemoved: { + addListener: vi.fn(), + }, + }, + storage: { + session: { + get: storageGetMock, + set: storageSetMock, + }, + }, + }, + configurable: true, + }); + }); + + it('restores enabled state after service worker restart', async () => { + const { StateManager } = await import('./state'); + const firstWorkerManager = new StateManager(); + await firstWorkerManager.setTabEnabled(tabId, true); + + const restartedWorkerManager = new StateManager(); + const restored = await restartedWorkerManager.getTabRuntimeState(tabId); + + expect(restored.enabled).toBe(true); + }); + + it('persists runtime state updates to chrome.storage.session', async () => { + const { StateManager } = await import('./state'); + const manager = new StateManager(); + await manager.setTabEnabled(tabId, true); + + expect(storageSetMock).toHaveBeenCalled(); + }); + + it('serializes writes so stale snapshots cannot overwrite newer state', async () => { + const pendingWrites: Array<{ + items: Record; + resolve: () => void; + }> = []; + + storageSetMock.mockImplementation((items: Record) => { + const snapshot = JSON.parse(JSON.stringify(items)) as Record; + return new Promise((resolve) => { + pendingWrites.push({ + items: snapshot, + resolve: () => { + Object.assign(storageBucket, snapshot); + resolve(); + }, + }); + }); + }); + + const { StateManager } = await import('./state'); + const manager = new StateManager(); + + const firstUpdate = manager.setTabEnabled(tabId, true); + const secondUpdate = manager.setTabEnabled(tabId, false); + + await vi.waitFor(() => { + expect(storageSetMock).toHaveBeenCalledTimes(1); + }); + expect(storageSetMock).toHaveBeenCalledTimes(1); + expect(pendingWrites).toHaveLength(1); + + const firstPendingWrite = pendingWrites[0]; + if (!firstPendingWrite) { + throw new Error('Expected first pending write to exist'); + } + firstPendingWrite.resolve(); + + await vi.waitFor(() => { + expect(storageSetMock).toHaveBeenCalledTimes(2); + }); + expect(storageSetMock).toHaveBeenCalledTimes(2); + expect(pendingWrites).toHaveLength(2); + + const secondPendingWrite = pendingWrites[1]; + if (!secondPendingWrite) { + throw new Error('Expected second pending write to exist'); + } + secondPendingWrite.resolve(); + await Promise.all([firstUpdate, secondUpdate]); + + const restartedWorkerManager = new StateManager(); + const restored = await restartedWorkerManager.getTabRuntimeState(tabId); + expect(restored.enabled).toBe(false); + }); +}); diff --git a/packages/extension/src/background/state.ts b/packages/extension/src/background/state.ts index f6ea6f6..b3754f2 100644 --- a/packages/extension/src/background/state.ts +++ b/packages/extension/src/background/state.ts @@ -2,24 +2,138 @@ * Per-tab state manager * Tracks tab runtime state (enabled + annotate mode) for each tab */ +const TAB_RUNTIME_STATES_STORAGE_KEY = 'onui_tab_runtime_states'; +const LOG_PREFIX = '[onUI][background][state]'; + +type TabRuntimeState = { enabled: boolean; annotateMode: boolean }; + export class StateManager { - private tabStates: Map = new Map(); + private tabStates: Map = new Map(); + private hydrationPromise: Promise | null = null; + private persistQueue: Promise = Promise.resolve(); constructor() { // Clean up state when tabs are closed chrome.tabs.onRemoved.addListener((tabId) => { - this.tabStates.delete(tabId); + void this.clearTabState(tabId); }); } - private getDefaultState(): { enabled: boolean; annotateMode: boolean } { + private getDefaultState(): TabRuntimeState { return { enabled: false, annotateMode: false }; } + private getStorageArea(): chrome.storage.StorageArea | null { + if (typeof chrome === 'undefined' || !chrome.storage) { + return null; + } + + const storageWithSession = chrome.storage as typeof chrome.storage & { + session?: chrome.storage.StorageArea; + }; + + return storageWithSession.session ?? null; + } + + private parseTabRuntimeState(value: unknown): TabRuntimeState | null { + if (!value || typeof value !== 'object') { + return null; + } + + const candidate = value as Record; + if (typeof candidate.enabled !== 'boolean' || typeof candidate.annotateMode !== 'boolean') { + return null; + } + + return { + enabled: candidate.enabled, + annotateMode: candidate.annotateMode, + }; + } + + private async ensureHydrated(): Promise { + if (this.hydrationPromise) { + return this.hydrationPromise; + } + + this.hydrationPromise = (async () => { + const storageArea = this.getStorageArea(); + if (!storageArea) { + return; + } + + try { + const result = await storageArea.get(TAB_RUNTIME_STATES_STORAGE_KEY); + const rawStates = result[TAB_RUNTIME_STATES_STORAGE_KEY]; + if (!rawStates || typeof rawStates !== 'object') { + return; + } + + for (const [tabIdString, rawState] of Object.entries(rawStates as Record)) { + const tabId = Number(tabIdString); + if (!Number.isFinite(tabId)) { + continue; + } + + const parsedState = this.parseTabRuntimeState(rawState); + if (!parsedState) { + continue; + } + + this.tabStates.set(tabId, parsedState); + } + } catch (error) { + console.warn(`${LOG_PREFIX} Failed to hydrate runtime state from session storage`, { + error: error instanceof Error ? error.message : 'unknown', + }); + } + })(); + + return this.hydrationPromise; + } + + private persistStates(): Promise { + this.persistQueue = this.persistQueue + .catch(() => { + // Keep the queue healthy even if a previous write failed. + }) + .then(async () => { + const storageArea = this.getStorageArea(); + if (!storageArea) { + return; + } + + const serialized = Array.from(this.tabStates.entries()).reduce>( + (acc, [tabId, state]) => { + acc[String(tabId)] = state; + return acc; + }, + {} + ); + + try { + await storageArea.set({ [TAB_RUNTIME_STATES_STORAGE_KEY]: serialized }); + } catch (error) { + console.warn(`${LOG_PREFIX} Failed to persist runtime state to session storage`, { + error: error instanceof Error ? error.message : 'unknown', + }); + } + }); + + return this.persistQueue; + } + + private async clearTabState(tabId: number): Promise { + await this.ensureHydrated(); + this.tabStates.delete(tabId); + await this.persistStates(); + } + /** * Get runtime state for a tab */ - getTabRuntimeState(tabId: number): { enabled: boolean; annotateMode: boolean } { + async getTabRuntimeState(tabId: number): Promise { + await this.ensureHydrated(); return this.tabStates.get(tabId) ?? this.getDefaultState(); } @@ -27,13 +141,14 @@ export class StateManager { * Set onUI enabled/disabled state for a tab. * Disabling always clears annotate mode. */ - setTabEnabled(tabId: number, enabled: boolean): { enabled: boolean; annotateMode: boolean } { - const current = this.getTabRuntimeState(tabId); + async setTabEnabled(tabId: number, enabled: boolean): Promise { + const current = await this.getTabRuntimeState(tabId); const next = { enabled, annotateMode: enabled ? current.annotateMode : false, }; this.tabStates.set(tabId, next); + await this.persistStates(); return next; } @@ -41,13 +156,14 @@ export class StateManager { * Set annotate mode for a tab. * Guard: annotate mode cannot be true when tab is disabled. */ - setAnnotateMode(tabId: number, annotateMode: boolean): { enabled: boolean; annotateMode: boolean } { - const current = this.getTabRuntimeState(tabId); + async setAnnotateMode(tabId: number, annotateMode: boolean): Promise { + const current = await this.getTabRuntimeState(tabId); const next = { enabled: current.enabled, annotateMode: current.enabled ? annotateMode : false, }; this.tabStates.set(tabId, next); + await this.persistStates(); return next; } @@ -55,24 +171,25 @@ export class StateManager { * Compatibility getter for legacy active state API. * Maps directly to annotate mode. */ - getState(tabId: number): boolean { - return this.getTabRuntimeState(tabId).annotateMode; + async getState(tabId: number): Promise { + const state = await this.getTabRuntimeState(tabId); + return state.annotateMode; } /** * Compatibility setter for legacy active state API. * Delegates to annotate mode setter. */ - setState(tabId: number, isActive: boolean): void { - this.setAnnotateMode(tabId, isActive); + async setState(tabId: number, isActive: boolean): Promise { + await this.setAnnotateMode(tabId, isActive); } /** * Compatibility toggle for legacy active state API. */ - toggleState(tabId: number): boolean { - const current = this.getState(tabId); - const next = this.setAnnotateMode(tabId, !current); + async toggleState(tabId: number): Promise { + const current = await this.getState(tabId); + const next = await this.setAnnotateMode(tabId, !current); return next.annotateMode; } } From 00aa82cdd6a6996bdeea012497f2f1cae7745cf2 Mon Sep 17 00:00:00 2001 From: Steve Olson Date: Sat, 28 Feb 2026 15:19:12 -0500 Subject: [PATCH 2/2] chore(extension): suppress onUI debug console.log noise by default Add shared log filter that suppresses [onUI...] console.log output in production builds while keeping warnings/errors. Initialize filter in content, background, and popup entrypoints. Keep logs opt-in via localStorage key onui_debug_logs=1 or true. --- packages/extension/src/background/index.ts | 3 ++ packages/extension/src/content/index.ts | 3 ++ packages/extension/src/popup/index.tsx | 4 ++ packages/extension/src/shared/logging.ts | 46 ++++++++++++++++++++++ 4 files changed, 56 insertions(+) create mode 100644 packages/extension/src/shared/logging.ts diff --git a/packages/extension/src/background/index.ts b/packages/extension/src/background/index.ts index 37ad6ad..c58b23e 100644 --- a/packages/extension/src/background/index.ts +++ b/packages/extension/src/background/index.ts @@ -4,7 +4,10 @@ import { setupMessageHandler } from './messages'; import { bootstrapNativeSync } from './native-sync'; +import { suppressOnUiDebugLogs } from '@/shared/logging'; +// Keep onUI debug logs opt-in in production builds. +suppressOnUiDebugLogs(); console.log('[onUI] Background service worker initialized'); // Set up message handling diff --git a/packages/extension/src/content/index.ts b/packages/extension/src/content/index.ts index 0b73cda..2e8db8d 100644 --- a/packages/extension/src/content/index.ts +++ b/packages/extension/src/content/index.ts @@ -3,7 +3,10 @@ */ import { renderApp } from './render'; +import { suppressOnUiDebugLogs } from '@/shared/logging'; +// Keep onUI debug logs opt-in in production builds. +suppressOnUiDebugLogs(); console.log('[onUI] Content script loaded'); declare global { diff --git a/packages/extension/src/popup/index.tsx b/packages/extension/src/popup/index.tsx index 86237a6..c33d05b 100644 --- a/packages/extension/src/popup/index.tsx +++ b/packages/extension/src/popup/index.tsx @@ -1,6 +1,10 @@ import { render } from 'preact'; import { Popup } from './components/Popup'; import './styles/popup.css'; +import { suppressOnUiDebugLogs } from '@/shared/logging'; + +// Keep onUI debug logs opt-in in production builds. +suppressOnUiDebugLogs(); const container = document.getElementById('app'); if (container) { diff --git a/packages/extension/src/shared/logging.ts b/packages/extension/src/shared/logging.ts new file mode 100644 index 0000000..6f0524b --- /dev/null +++ b/packages/extension/src/shared/logging.ts @@ -0,0 +1,46 @@ +const DEBUG_LOG_STORAGE_KEY = 'onui_debug_logs'; +const PATCH_FLAG = '__onui_debug_log_filter_installed__'; + +type ConsoleWithPatchFlag = Console & { + [PATCH_FLAG]?: boolean; +}; + +function isDebugLoggingEnabled(): boolean { + if (import.meta.env.DEV) { + return true; + } + + try { + if (typeof localStorage === 'undefined') { + return false; + } + + const value = localStorage.getItem(DEBUG_LOG_STORAGE_KEY); + return value === '1' || value === 'true'; + } catch { + return false; + } +} + +export function suppressOnUiDebugLogs(): void { + if (isDebugLoggingEnabled()) { + return; + } + + const consoleWithPatchFlag = console as ConsoleWithPatchFlag; + if (consoleWithPatchFlag[PATCH_FLAG]) { + return; + } + + const originalLog = console.log.bind(console); + console.log = (...args: unknown[]) => { + const first = args[0]; + if (typeof first === 'string' && first.startsWith('[onUI')) { + return; + } + + originalLog(...args); + }; + + consoleWithPatchFlag[PATCH_FLAG] = true; +}