From 324cfcc4865ca5e288e91b4f5d34fa411eb7e6c1 Mon Sep 17 00:00:00 2001 From: john kevin asprec Date: Tue, 14 Jul 2026 21:33:56 +0800 Subject: [PATCH] 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 --- entrypoints/background.ts | 22 +---- entrypoints/content.ts | 38 +++----- entrypoints/options/Options.tsx | 35 +------- entrypoints/popup/Popup.tsx | 47 ++-------- src/lib/messaging.ts | 68 +++++++++++++++ src/lib/types.ts | 54 ++++++++++++ tests/unit/messaging.test.ts | 149 ++++++++++++++++++++++++++++++++ tsconfig.json | 2 +- vitest.config.ts | 6 ++ wxt.config.ts | 6 ++ 10 files changed, 303 insertions(+), 124 deletions(-) create mode 100644 src/lib/messaging.ts create mode 100644 src/lib/types.ts create mode 100644 tests/unit/messaging.test.ts diff --git a/entrypoints/background.ts b/entrypoints/background.ts index eaee598..14563dc 100644 --- a/entrypoints/background.ts +++ b/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 { diff --git a/entrypoints/content.ts b/entrypoints/content.ts index 83ae6ec..4423942 100644 --- a/entrypoints/content.ts +++ b/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: [''], @@ -362,15 +363,7 @@ export default defineContentScript({ } } - // ─── Extension context guard ────────────────────────────────────────────── - - function isExtensionValid(): boolean { - try { - return typeof chrome !== 'undefined' && !!chrome.runtime?.id; - } catch { - return false; - } - } + // Extension context guard: isExtensionValid is imported from ~/lib/messaging. 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): Promise { - if (!isExtensionValid()) { + async function sendToBackground(payload: Record): Promise { + 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; - } + hideToolbar(); + }); } // ─── 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; diff --git a/entrypoints/options/Options.tsx b/entrypoints/options/Options.tsx index 69db6ff..d5feb3e 100644 --- a/entrypoints/options/Options.tsx +++ b/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) => 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, 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): Promise { - 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() { diff --git a/entrypoints/popup/Popup.tsx b/entrypoints/popup/Popup.tsx index 3df018a..cda530b 100644 --- a/entrypoints/popup/Popup.tsx +++ b/entrypoints/popup/Popup.tsx @@ -1,41 +1,6 @@ import React, { useEffect, useRef, useState } from 'react'; import { createRoot } from 'react-dom/client'; - -// ─── Safe chrome storage wrapper ───────────────────────────────────────────── - -function safeStorageGet( - storage: chrome.storage.StorageArea, - keys: string[], - callback: (result: Record) => 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) { - 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, -): Promise | 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; - } -} +import { safeSendMessage, safeStorageGet, safeStorageSet } from '@lib/messaging'; // ─── 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) => { 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) => { const val = e.target.value; setWritingStyle(val); - safeStorageSet(chrome.storage.local, { writingStyle: val }); + safeStorageSet({ writingStyle: val }); }; const openSettings = () => { diff --git a/src/lib/messaging.ts b/src/lib/messaging.ts new file mode 100644 index 0000000..a346570 --- /dev/null +++ b/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) => 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, + 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( + message: Record, + onContextInvalidated?: () => void, +): Promise { + 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; + } +} diff --git a/src/lib/types.ts b/src/lib/types.ts new file mode 100644 index 0000000..93e7cd0 --- /dev/null +++ b/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; diff --git a/tests/unit/messaging.test.ts b/tests/unit/messaging.test.ts new file mode 100644 index 0000000..6df7f4f --- /dev/null +++ b/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) => void) => cb({})), + set: vi.fn((_data: Record, cb?: () => void) => cb && cb()), + }; + const session = { + get: vi.fn((_keys: string[], cb: (r: Record) => void) => cb({})), + set: vi.fn((_data: Record, 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(); + }); +}); diff --git a/tsconfig.json b/tsconfig.json index 332b67d..52c16cf 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -13,7 +13,7 @@ "resolveJsonModule": true, "allowImportingTsExtensions": true, "paths": { - "~/*": ["./src/*"] + "@lib/*": ["./src/lib/*"] } }, "include": [ diff --git a/vitest.config.ts b/vitest.config.ts index dd0d6b9..a176bfb 100644 --- a/vitest.config.ts +++ b/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, diff --git a/wxt.config.ts b/wxt.config.ts index 8038b04..fc359ea 100644 --- a/wxt.config.ts +++ b/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',