fix(extension): persist and serialize per-tab runtime state

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

packages/extension/src/background/messages.ts

@@ -25,7 +25,7 @@ async function ensureContentScriptInjected(tabId: number): Promise<void> {
 }
 
 async function reInjectContentScriptIfEnabled(tabId: number, url?: string): Promise<void> {
-  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<void> {
-  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,

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<string, unknown>;
+  let storageSetMock: ReturnType<typeof vi.fn>;
+
+  beforeEach(() => {
+    vi.resetModules();
+    storageBucket = {};
+    storageSetMock = vi.fn(async (items: Record<string, unknown>) => {
+      Object.assign(storageBucket, items);
+    });
+
+    const storageGetMock = vi.fn(async (keys?: string | string[] | Record<string, unknown>) => {
+      if (typeof keys === 'string') {
+        return { [keys]: storageBucket[keys] };
+      }
+
+      if (Array.isArray(keys)) {
+        return keys.reduce<Record<string, unknown>>((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<string, unknown>;
+      resolve: () => void;
+    }> = [];
+
+    storageSetMock.mockImplementation((items: Record<string, unknown>) => {
+      const snapshot = JSON.parse(JSON.stringify(items)) as Record<string, unknown>;
+      return new Promise<void>((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);
+  });
+});