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/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; } } 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; +}