refactor: extract shared types + safe chrome wrappers into src/lib

- src/lib/types.ts: message contract (both ANALYZE_TEXT shapes preserved),
  config/response types, storage-key constants.
- src/lib/messaging.ts: single safeStorageGet/safeStorageSet/safeSendMessage/
  isExtensionValid implementation replacing the three divergent copies in
  Options, Popup, and content. Content script keeps its refresh-toast
  behavior via an onContextInvalidated callback.
- New '@lib' import alias (wxt force-overwrites '~' and '@' to srcDir, so
  those cannot point at ./src); wired in wxt.config, tsconfig, vitest.
- tests/unit/messaging.test.ts: 14 unit tests over the wrappers.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

entrypoints/background.ts

@@ -1,5 +1,6 @@
 import { defineBackground } from 'wxt/utils/define-background';
 import nacl from 'tweetnacl';
+import type { AnalyzePayload, LexAIConfig, LexAIResponse } from '@lib/types';
 
 // ─── Fetch with timeout ───────────────────────────────────────────────────────
 
@@ -13,27 +14,6 @@ async function fetchWithTimeout(url: string, options: RequestInit, timeoutMs = 3
   }
 }
 
-// ─── Types ────────────────────────────────────────────────────────────────────
-
-interface AnalyzePayload {
-  text: string;
-  action: string;
-  style?: string;
-}
-
-interface LexAIConfig {
-  provider?: string;
-  apiKey?: string;
-  apiKeyEnc?: string;
-  encKey?: string;
-  model?: string;
-}
-
-interface LexAIResponse {
-  result?: string;
-  error?: string;
-}
-
 // ─── System prompts ───────────────────────────────────────────────────────────
 
 function getSystemPrompt(action: string, style?: string): string {

entrypoints/content.ts

@@ -1,4 +1,5 @@
 import { defineContentScript } from 'wxt/utils/define-content-script';
+import { isExtensionValid, safeSendMessage } from '@lib/messaging';
 
 export default defineContentScript({
   matches: ['<all_urls>'],
@@ -362,15 +363,7 @@ export default defineContentScript({
       }
     }
 
-    // ─── Extension context guard ──────────────────────────────────────────────
+    // Extension context guard: isExtensionValid is imported from ~/lib/messaging.
-
-    function isExtensionValid(): boolean {
-      try {
-        return typeof chrome !== 'undefined' && !!chrome.runtime?.id;
-      } catch {
-        return false;
-      }
-    }
 
     function showErrorToast(msg: string) {
       const toast = document.createElement('div');
@@ -528,7 +521,7 @@ export default defineContentScript({
           grid.style.display       = 'none';
           spinnerSlot.style.display = 'flex';
 
-          const response = await safeSendMessage({
+          const response = await sendToBackground({
             type:   'COPY_AS',
             text,
             format: fmt,
@@ -563,22 +556,13 @@ export default defineContentScript({
     }
 
     // ─── Safe chrome.runtime.sendMessage wrapper ──────────────────────────────
+    // Shared implementation; on a stale extension context we toast and clean up.
 
-    async function safeSendMessage(payload: Record<string, unknown>): Promise<unknown> {
+    async function sendToBackground(payload: Record<string, unknown>): Promise<unknown> {
-      if (!isExtensionValid()) {
+      return safeSendMessage(payload, () => {
-        showErrorToast('LexAI was updated — please refresh this page.');
-        return null;
-      }
-      try {
-        return await chrome.runtime.sendMessage(payload);
-      } catch (err) {
-        if (String(err).includes('Extension context invalidated') ||
-            String(err).includes('message channel closed')) {
         showErrorToast('LexAI was updated — please refresh this page.');
         hideToolbar();
-        }
+      });
-        return null;
-      }
     }
 
     // ─── LLM call ─────────────────────────────────────────────────────────────
@@ -628,14 +612,14 @@ export default defineContentScript({
       }
 
       try {
-        const response = await safeSendMessage({
+        const response = await sendToBackground({
           type:    'ANALYZE_TEXT',
           payload: { text: textToProcess, action, style: currentStyle },
         }) as { error?: string; result?: string } | null;
 
         hideToolbar();
 
-        if (response === null) return; // safeSendMessage already handled the error
+        if (response === null) return; // sendToBackground already handled the error
 
         if (response?.error) {
           showModal(`❌ ${response.error}`, null, snapStart, snapEnd, snapElement, snapRange, action, textToProcess);
@@ -832,7 +816,7 @@ export default defineContentScript({
         regenBtn.disabled = true;
         styleSelect.disabled = true;
 
-        const response = await safeSendMessage({
+        const response = await sendToBackground({
           type:    'ANALYZE_TEXT',
           payload: { text: textForRegenerate, action, style: chosenStyle },
         }) as { error?: string; result?: string } | null;

entrypoints/options/Options.tsx

@@ -1,6 +1,7 @@
 import React, { useEffect, useRef, useState } from 'react';
 import { createRoot } from 'react-dom/client';
 import nacl from 'tweetnacl';
+import { safeSendMessage, safeStorageGet, safeStorageSet } from '@lib/messaging';
 
 // ─── Provider config ──────────────────────────────────────────────────────────
 
@@ -166,40 +167,6 @@ const styles = {
   } as React.CSSProperties,
 };
 
-// ─── Safe chrome storage wrappers ────────────────────────────────────────────
-
-function safeStorageGet(keys: string[], callback: (result: Record<string, string>) => void) {
-  try {
-    if (typeof chrome === 'undefined' || !chrome.runtime?.id) return;
-    chrome.storage.local.get(keys, callback);
-  } catch (err) {
-    console.warn('LexAI: Extension context invalidated', err);
-  }
-}
-
-function safeStorageSet(data: Record<string, string>, callback?: () => void) {
-  try {
-    if (typeof chrome === 'undefined' || !chrome.runtime?.id) return;
-    if (callback) {
-      chrome.storage.local.set(data, callback);
-    } else {
-      chrome.storage.local.set(data);
-    }
-  } catch (err) {
-    console.warn('LexAI: Extension context invalidated', err);
-  }
-}
-
-async function safeSendMessage(message: Record<string, unknown>): Promise<any> {
-  try {
-    if (typeof chrome === 'undefined' || !chrome.runtime?.id) return null;
-    return await chrome.runtime.sendMessage(message);
-  } catch (err) {
-    console.warn('LexAI: sendMessage failed', err);
-    return null;
-  }
-}
-
 // ─── Component ────────────────────────────────────────────────────────────────
 
 function OptionsPage() {

entrypoints/popup/Popup.tsx

@@ -1,41 +1,6 @@
 import React, { useEffect, useRef, useState } from 'react';
 import { createRoot } from 'react-dom/client';
-
+import { safeSendMessage, safeStorageGet, safeStorageSet } from '@lib/messaging';
-// ─── Safe chrome storage wrapper ─────────────────────────────────────────────
-
-function safeStorageGet(
-  storage: chrome.storage.StorageArea,
-  keys: string[],
-  callback: (result: Record<string, string>) => void,
-) {
-  try {
-    if (typeof chrome === 'undefined' || !chrome.runtime?.id) return;
-    storage.get(keys, callback);
-  } catch (err) {
-    console.warn('LexAI: Extension context invalidated', err);
-  }
-}
-
-function safeStorageSet(storage: chrome.storage.StorageArea, data: Record<string, string>) {
-  try {
-    if (typeof chrome === 'undefined' || !chrome.runtime?.id) return;
-    storage.set(data);
-  } catch (err) {
-    console.warn('LexAI: Extension context invalidated', err);
-  }
-}
-
-async function safeSendMessage(
-  message: Record<string, unknown>,
-): Promise<Record<string, string> | null> {
-  try {
-    if (typeof chrome === 'undefined' || !chrome.runtime?.id) return null;
-    return await chrome.runtime.sendMessage(message);
-  } catch (err) {
-    console.warn('LexAI: sendMessage failed', err);
-    return null;
-  }
-}
 
 // ─── Constants ───────────────────────────────────────────────────────────────
 
@@ -222,7 +187,7 @@ function Popup() {
 
   // Load config + restore session input + load writing style
   useEffect(() => {
-    safeStorageGet(chrome.storage.local, ['provider', 'apiKey', 'apiKeyEnc', 'model', 'writingStyle'], (result) => {
+    safeStorageGet(['provider', 'apiKey', 'apiKeyEnc', 'model', 'writingStyle'], (result) => {
       if (result.apiKey || result.apiKeyEnc) {
         setConfigured(true);
         setProvider(result.provider || 'openai');
@@ -235,11 +200,11 @@ function Popup() {
 
     if (!sessionRestored.current) {
       sessionRestored.current = true;
-      safeStorageGet(chrome.storage.session, ['lexai_popup_input'], (res) => {
+      safeStorageGet(['lexai_popup_input'], (res) => {
         if (res.lexai_popup_input) {
           setInputText(res.lexai_popup_input);
         }
-      });
+      }, chrome.storage.session);
     }
   }, []);
 
@@ -247,14 +212,14 @@ function Popup() {
   const handleInputChange = (e: React.ChangeEvent<HTMLTextAreaElement>) => {
     const val = e.target.value;
     setInputText(val);
-    safeStorageSet(chrome.storage.session, { lexai_popup_input: val });
+    safeStorageSet({ lexai_popup_input: val }, undefined, chrome.storage.session);
   };
 
   // Save writing style on change
   const handleStyleChange = (e: React.ChangeEvent<HTMLSelectElement>) => {
     const val = e.target.value;
     setWritingStyle(val);
-    safeStorageSet(chrome.storage.local, { writingStyle: val });
+    safeStorageSet({ writingStyle: val });
   };
 
   const openSettings = () => {

src/lib/messaging.ts

@@ -0,0 +1,68 @@
+// Single implementation of the "safe" chrome API wrappers shared by every
+// context. All of them guard against the extension context being invalidated
+// (extension reloaded/updated while a page, popup, or options view is open),
+// which otherwise throws from any chrome.* call.
+
+export function isExtensionValid(): boolean {
+  try {
+    return typeof chrome !== 'undefined' && !!chrome.runtime?.id;
+  } catch {
+    return false;
+  }
+}
+
+export function safeStorageGet(
+  keys: string[],
+  callback: (result: Record<string, any>) => void,
+  storage: chrome.storage.StorageArea = chrome.storage.local,
+): void {
+  try {
+    if (!isExtensionValid()) return;
+    storage.get(keys, callback);
+  } catch (err) {
+    console.warn('LexAI: Extension context invalidated', err);
+  }
+}
+
+export function safeStorageSet(
+  data: Record<string, unknown>,
+  callback?: () => void,
+  storage: chrome.storage.StorageArea = chrome.storage.local,
+): void {
+  try {
+    if (!isExtensionValid()) return;
+    if (callback) {
+      storage.set(data, callback);
+    } else {
+      storage.set(data);
+    }
+  } catch (err) {
+    console.warn('LexAI: Extension context invalidated', err);
+  }
+}
+
+// Sends a message to the background worker. Returns null (never throws) when
+// the extension context is gone or the send fails. `onContextInvalidated` lets
+// UI contexts react (e.g. the content script shows a "please refresh" toast).
+export async function safeSendMessage<T = any>(
+  message: Record<string, unknown>,
+  onContextInvalidated?: () => void,
+): Promise<T | null> {
+  if (!isExtensionValid()) {
+    onContextInvalidated?.();
+    return null;
+  }
+  try {
+    return await chrome.runtime.sendMessage(message);
+  } catch (err) {
+    if (
+      String(err).includes('Extension context invalidated') ||
+      String(err).includes('message channel closed')
+    ) {
+      onContextInvalidated?.();
+    } else {
+      console.warn('LexAI: sendMessage failed', err);
+    }
+    return null;
+  }
+}

src/lib/types.ts

@@ -0,0 +1,54 @@
+// Shared types and constants for all extension contexts (background, content,
+// options, popup). This is the single source of truth for the message contract
+// and the storage schema.
+
+export interface AnalyzePayload {
+  text: string;
+  action: string;
+  style?: string;
+}
+
+export interface LexAIConfig {
+  provider?: string;
+  apiKey?: string; // legacy plaintext key — kept for backward compat until migrated
+  apiKeyEnc?: string; // base64(nonce + secretbox ciphertext)
+  encKey?: string; // base64 32-byte secretbox key
+  model?: string;
+}
+
+export interface LexAIResponse {
+  result?: string;
+  error?: string;
+}
+
+// Storage keys the background worker reads when resolving provider config.
+export const CONFIG_STORAGE_KEYS = ['provider', 'apiKey', 'apiKeyEnc', 'encKey', 'model'] as const;
+
+// ─── Message contract ─────────────────────────────────────────────────────────
+// ANALYZE_TEXT intentionally supports BOTH shapes:
+//   { type, payload: { text, action, style } }  (content script, popup)
+//   { type, text, action, style }               (flat/legacy)
+// The background handler normalizes; both must keep working.
+
+export interface AnalyzeTextMessage {
+  type: 'ANALYZE_TEXT';
+  payload?: AnalyzePayload;
+  text?: string;
+  action?: string;
+  style?: string;
+}
+
+export interface CopyAsMessage {
+  type: 'COPY_AS';
+  text: string;
+  format: string;
+}
+
+export interface ListModelsMessage {
+  type: 'LIST_MODELS';
+  provider?: string;
+  // Freshly typed key not yet saved — preferred over the stored key when present.
+  apiKey?: string;
+}
+
+export type LexAIMessage = AnalyzeTextMessage | CopyAsMessage | ListModelsMessage;

tests/unit/messaging.test.ts

@@ -0,0 +1,149 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import {
+  isExtensionValid,
+  safeStorageGet,
+  safeStorageSet,
+  safeSendMessage,
+} from '@lib/messaging';
+
+// Rebuild the chrome mock per test so each case controls runtime.id presence.
+function mockChrome({ valid = true }: { valid?: boolean } = {}) {
+  const local = {
+    get: vi.fn((_keys: string[], cb: (r: Record<string, unknown>) => void) => cb({})),
+    set: vi.fn((_data: Record<string, unknown>, cb?: () => void) => cb && cb()),
+  };
+  const session = {
+    get: vi.fn((_keys: string[], cb: (r: Record<string, unknown>) => void) => cb({})),
+    set: vi.fn((_data: Record<string, unknown>, cb?: () => void) => cb && cb()),
+  };
+  global.chrome = {
+    storage: { local, session },
+    runtime: {
+      id: valid ? 'test-extension-id' : undefined,
+      sendMessage: vi.fn().mockResolvedValue({ result: 'ok' }),
+    },
+  } as any;
+  return { local, session };
+}
+
+beforeEach(() => {
+  vi.restoreAllMocks();
+});
+
+describe('isExtensionValid', () => {
+  it('is true when chrome.runtime.id exists', () => {
+    mockChrome({ valid: true });
+    expect(isExtensionValid()).toBe(true);
+  });
+
+  it('is false when the extension context is invalidated (no runtime.id)', () => {
+    mockChrome({ valid: false });
+    expect(isExtensionValid()).toBe(false);
+  });
+});
+
+describe('safeStorageGet', () => {
+  it('reads from chrome.storage.local by default', () => {
+    const { local } = mockChrome();
+    const cb = vi.fn();
+    safeStorageGet(['provider', 'model'], cb);
+    expect(local.get).toHaveBeenCalledWith(['provider', 'model'], cb);
+    expect(cb).toHaveBeenCalledWith({});
+  });
+
+  it('reads from an explicitly passed storage area (e.g. session)', () => {
+    const { local, session } = mockChrome();
+    const cb = vi.fn();
+    safeStorageGet(['lexai_popup_input'], cb, (global.chrome as any).storage.session);
+    expect(session.get).toHaveBeenCalled();
+    expect(local.get).not.toHaveBeenCalled();
+  });
+
+  it('is a silent no-op when the extension context is gone', () => {
+    const { local } = mockChrome({ valid: false });
+    const cb = vi.fn();
+    safeStorageGet(['provider'], cb);
+    expect(local.get).not.toHaveBeenCalled();
+    expect(cb).not.toHaveBeenCalled();
+  });
+
+  it('swallows a throwing storage call with a warning', () => {
+    mockChrome();
+    (global.chrome as any).storage.local.get = vi.fn(() => {
+      throw new Error('Extension context invalidated.');
+    });
+    const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
+    expect(() => safeStorageGet(['provider'], vi.fn())).not.toThrow();
+    expect(warn).toHaveBeenCalled();
+  });
+});
+
+describe('safeStorageSet', () => {
+  it('writes to chrome.storage.local by default and honors the callback', () => {
+    const { local } = mockChrome();
+    const done = vi.fn();
+    safeStorageSet({ provider: 'openai' }, done);
+    expect(local.set).toHaveBeenCalledWith({ provider: 'openai' }, done);
+    expect(done).toHaveBeenCalled();
+  });
+
+  it('writes without a callback', () => {
+    const { local } = mockChrome();
+    safeStorageSet({ writingStyle: 'Formal' });
+    expect(local.set).toHaveBeenCalledWith({ writingStyle: 'Formal' });
+  });
+
+  it('writes to an explicitly passed storage area', () => {
+    const { session } = mockChrome();
+    safeStorageSet({ lexai_popup_input: 'draft' }, undefined, (global.chrome as any).storage.session);
+    expect(session.set).toHaveBeenCalledWith({ lexai_popup_input: 'draft' });
+  });
+
+  it('is a no-op when the extension context is gone', () => {
+    const { local } = mockChrome({ valid: false });
+    safeStorageSet({ provider: 'openai' });
+    expect(local.set).not.toHaveBeenCalled();
+  });
+});
+
+describe('safeSendMessage', () => {
+  it('resolves with the background response', async () => {
+    mockChrome();
+    const res = await safeSendMessage({ type: 'ANALYZE_TEXT', payload: { text: 'hi', action: 'grammar' } });
+    expect(res).toEqual({ result: 'ok' });
+    expect((global.chrome as any).runtime.sendMessage).toHaveBeenCalledWith({
+      type: 'ANALYZE_TEXT',
+      payload: { text: 'hi', action: 'grammar' },
+    });
+  });
+
+  it('returns null and fires onContextInvalidated when the context is gone up front', async () => {
+    mockChrome({ valid: false });
+    const onInvalid = vi.fn();
+    const res = await safeSendMessage({ type: 'LIST_MODELS' }, onInvalid);
+    expect(res).toBeNull();
+    expect(onInvalid).toHaveBeenCalledTimes(1);
+  });
+
+  it('returns null and fires onContextInvalidated when sendMessage rejects with an invalidation error', async () => {
+    mockChrome();
+    (global.chrome as any).runtime.sendMessage = vi
+      .fn()
+      .mockRejectedValue(new Error('Extension context invalidated.'));
+    const onInvalid = vi.fn();
+    const res = await safeSendMessage({ type: 'COPY_AS', text: 'x', format: 'markdown' }, onInvalid);
+    expect(res).toBeNull();
+    expect(onInvalid).toHaveBeenCalledTimes(1);
+  });
+
+  it('returns null and warns (no invalidation callback) on other send failures', async () => {
+    mockChrome();
+    (global.chrome as any).runtime.sendMessage = vi.fn().mockRejectedValue(new Error('boom'));
+    const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
+    const onInvalid = vi.fn();
+    const res = await safeSendMessage({ type: 'LIST_MODELS' }, onInvalid);
+    expect(res).toBeNull();
+    expect(onInvalid).not.toHaveBeenCalled();
+    expect(warn).toHaveBeenCalled();
+  });
+});

tsconfig.json

@@ -13,7 +13,7 @@
     "resolveJsonModule": true,
     "allowImportingTsExtensions": true,
     "paths": {
-      "~/*": ["./src/*"]
+      "@lib/*": ["./src/lib/*"]
     }
   },
   "include": [

vitest.config.ts

@@ -1,6 +1,12 @@
+import { fileURLToPath } from 'node:url';
 import { defineConfig } from 'vitest/config';
 
 export default defineConfig({
+  resolve: {
+    alias: {
+      '@lib': fileURLToPath(new URL('./src/lib', import.meta.url)),
+    },
+  },
   test: {
     environment: 'jsdom',
     globals: true,

wxt.config.ts

@@ -1,8 +1,14 @@
+import { resolve } from 'node:path';
 import { defineConfig } from 'wxt';
 import pkg from './package.json';
 
 export default defineConfig({
   extensionApi: 'chrome',
+  // Shared-code alias. Deliberately NOT "~" or "@" — WXT force-overwrites those
+  // to srcDir (the project root here), so they cannot point at ./src.
+  alias: {
+    '@lib': resolve(__dirname, 'src/lib'),
+  },
   modules: ['@wxt-dev/module-react'],
   manifest: {
     name: 'LexAI - AI Writing Assistant',