feat: integrate storage with jotai (#5)

apps/ai-assistant/package.json

@@ -11,7 +11,9 @@
     "@douyinfe/semi-ui": "^2.51.3",
     "@google/generative-ai": "^0.1.3",
     "@webx-kit/runtime": "workspace:^",
+    "@webx-kit/storage": "workspace:^",
     "clsx": "^2.1.0",
+    "jotai": "^2.6.2",
     "react": "^18.2.0",
     "react-dom": "^18.2.0"
   },

apps/ai-assistant/src/content-scripts/app.tsx

@@ -10,6 +10,8 @@ import {
   rangeToReference,
 } from '@webx-kit/runtime/content-scripts';
 import clsx from 'clsx';
+import { atom, useAtomValue } from 'jotai';
+import { apiKeyAtom } from '@/hooks/atoms/config';
 import { Provider } from './features/provider';
 import './global.less';
 
@@ -25,6 +27,8 @@ export const App = () => {
   const isDarkMode = useMemo(isPageInDark, [visible]);
   const containerRef = useRef<HTMLDivElement>(null);
 
+  const genAI = useAtomValue(genAIAtom);
+
   useEffect(() => {
     const ac = new AbortController();
     document.addEventListener(
@@ -191,10 +195,7 @@ export const App = () => {
   );
 };
 
-let genAI: GoogleGenerativeAI | undefined;
-
-const GOOGLE_API_KEY = 'GOOGLE_API_KEY';
-chrome.storage.local.get(GOOGLE_API_KEY).then(({ GOOGLE_API_KEY }) => {
-  if (!GOOGLE_API_KEY) return console.warn('`GOOGLE_API_KEY` is not provided.`');
-  genAI = new GoogleGenerativeAI(GOOGLE_API_KEY);
+const genAIAtom = atom(async (get) => {
+  const apiKey = await get(apiKeyAtom);
+  return apiKey ? new GoogleGenerativeAI(apiKey) : null;
 });

apps/ai-assistant/src/hooks/atoms/config.ts

@@ -0,0 +1,6 @@
+import { atomWithStorage } from 'jotai/utils';
+import { createStorage } from '@webx-kit/storage';
+
+const chromeStorage = createStorage();
+
+export const apiKeyAtom = atomWithStorage<string | null>('apiKey', null, chromeStorage);

apps/ai-assistant/src/pages/options/app.tsx

@@ -1,5 +1,7 @@
 import { useEffect, useRef } from 'react';
 import { Button, Form, Toast } from '@douyinfe/semi-ui';
+import { useAtom } from 'jotai';
+import { apiKeyAtom } from '@/hooks/atoms/config';
 import { useSemiTheme } from '@/hooks/use-theme';
 
 interface FormValues {
@@ -9,15 +11,14 @@ interface FormValues {
 export const App = () => {
   useSemiTheme();
   const formRef = useRef<Form<FormValues>>(null);
+  const [apiKey, setAPIKey] = useAtom(apiKeyAtom);
 
   useEffect(() => {
-    chrome.storage.local.get('GOOGLE_API_KEY', ({ GOOGLE_API_KEY }) => {
-      formRef.current?.formApi.setValue('apiKey', GOOGLE_API_KEY || '');
-    });
-  }, []);
+    if (apiKey) formRef.current?.formApi.setValue('apiKey', apiKey);
+  }, [apiKey]);
 
   const handleSubmit = async (formValues: FormValues) => {
-    await chrome.storage.local.set({ GOOGLE_API_KEY: formValues.apiKey });
+    setAPIKey(formValues.apiKey);
     Toast.success('Saved');
   };
 

cspell.json

@@ -13,6 +13,7 @@
     "rsbuild",
     "svgr",
     "tailwindcss",
+    "testid",
     "webx"
   ]
 }

packages/storage/demo/content-scripts/index.tsx

@@ -0,0 +1,16 @@
+import { createRoot } from 'react-dom/client';
+import { useAtomValue } from 'jotai/react';
+import { apiKeyAtom } from '../hooks/atoms';
+
+const App = () => {
+  const apiKey = useAtomValue(apiKeyAtom);
+  return (
+    <div>
+      <div data-testid="apiKey">{apiKey}</div>
+    </div>
+  );
+};
+
+const root = document.createElement('div');
+createRoot(root).render(<App />);
+document.body.append(root);

packages/storage/demo/hooks/atoms.ts

@@ -0,0 +1,6 @@
+import { atomWithStorage } from 'jotai/utils';
+import { createStorage } from '@/index';
+
+const chromeStorage = createStorage();
+
+export const apiKeyAtom = atomWithStorage<string>('apiKey', 'DEFAULT', chromeStorage);

packages/storage/demo/manifest.ts

@@ -1,5 +1,3 @@
-const isDev = process.env.NODE_ENV === 'development';
-
 const manifest: chrome.runtime.ManifestV3 = {
   manifest_version: 3,
   name: 'WebX Kit Storage',
@@ -9,9 +7,6 @@ const manifest: chrome.runtime.ManifestV3 = {
     service_worker: 'static/js/background.js',
     type: 'module',
   },
-  action: {
-    default_popup: 'popup.html',
-  },
   options_ui: {
     page: 'options.html',
     open_in_tab: true,
@@ -24,14 +19,6 @@ const manifest: chrome.runtime.ManifestV3 = {
     },
   ],
   host_permissions: ['<all_urls>'],
-
-  ...(isDev
-    ? {
-        content_security_policy: {
-          extension_pages: `script-src 'self' http://localhost:${process.env.PORT}/; object-src 'self' http://localhost:${process.env.PORT}/`,
-        },
-      }
-    : {}),
 };
 
 export default manifest;

packages/storage/demo/pages/jotai/index.tsx

@@ -0,0 +1,22 @@
+import { createRoot } from 'react-dom/client';
+import { useAtom } from 'jotai/react';
+import { apiKeyAtom } from '../../hooks/atoms';
+
+const App = () => {
+  const [apiKey, setAPIKey] = useAtom(apiKeyAtom);
+  return (
+    <div>
+      <div data-testid="apiKey">{apiKey}</div>
+      <button
+        data-testid="change"
+        onClick={() => {
+          setAPIKey('Changed');
+        }}
+      >
+        Change
+      </button>
+    </div>
+  );
+};
+
+createRoot(document.getElementById('root')!).render(<App />);

packages/storage/demo/pages/options/index.tsx

@@ -0,0 +1,14 @@
+import { createRoot } from 'react-dom/client';
+import { useAtomValue } from 'jotai/react';
+import { apiKeyAtom } from '../../hooks/atoms';
+
+const App = () => {
+  const apiKey = useAtomValue(apiKeyAtom);
+  return (
+    <div>
+      <div data-testid="apiKey">{apiKey}</div>
+    </div>
+  );
+};
+
+createRoot(document.getElementById('root')!).render(<App />);

packages/storage/e2e/jotai.spec.ts

@@ -0,0 +1,39 @@
+import { setupStaticServer } from '@webx-kit/test-utils/playwright';
+import { expect, test } from './context';
+
+const getWebpageURL = setupStaticServer(test);
+
+test('Single page', async ({ getURL, page }) => {
+  await page.goto(await getURL('jotai.html'));
+
+  await expect(page.getByTestId('apiKey')).toHaveText('DEFAULT');
+  await page.getByTestId('change').click();
+  await expect(page.getByTestId('apiKey')).toHaveText('Changed');
+
+  await page.reload();
+  await expect(page.getByTestId('apiKey')).toHaveText('Changed');
+});
+
+test('Cross pages', async ({ getURL, context }) => {
+  const jotaiPage = await context.newPage();
+  const optionsPage = await context.newPage();
+  const webPage = await context.newPage();
+
+  await jotaiPage.goto(await getURL('jotai.html'));
+  await optionsPage.goto(await getURL('options.html'));
+  await webPage.goto(getWebpageURL());
+
+  await expect(jotaiPage.getByTestId('apiKey')).toHaveText('DEFAULT');
+  await expect(optionsPage.getByTestId('apiKey')).toHaveText('DEFAULT');
+  await expect(webPage.getByTestId('apiKey')).toHaveText('DEFAULT');
+
+  await jotaiPage.getByTestId('change').click();
+  await expect(jotaiPage.getByTestId('apiKey')).toHaveText('Changed');
+  await expect(optionsPage.getByTestId('apiKey')).toHaveText('Changed');
+  await expect(webPage.getByTestId('apiKey')).toHaveText('Changed');
+
+  await Promise.all([jotaiPage.reload(), optionsPage.reload(), webPage.reload()]);
+  await expect(jotaiPage.getByTestId('apiKey')).toHaveText('Changed');
+  await expect(optionsPage.getByTestId('apiKey')).toHaveText('Changed');
+  await expect(webPage.getByTestId('apiKey')).toHaveText('Changed');
+});

packages/storage/modern.config.ts

@@ -7,7 +7,7 @@ export default defineConfig({
     appTools(),
     webxPlugin({
       background: './demo/background/index.ts',
-      contentScripts: './demo/content-scripts/index.ts',
+      contentScripts: './demo/content-scripts/index.tsx',
       manifest: './demo/manifest.ts',
     }),
   ],

packages/storage/package.json

@@ -1,20 +1,48 @@
 {
   "name": "@webx-kit/storage",
   "version": "0.0.0",
+  "exports": {
+    ".": {
+      "types": "./src/index.ts",
+      "default": "./src/index.ts"
+    }
+  },
+  "types": "./src/index.ts",
+  "publishConfig": {
+    "exports": {
+      ".": {
+        "types": "./dist/index.d.mts",
+        "default": "./dist/index.mjs"
+      }
+    },
+    "types": "./dist/index.d.mts"
+  },
+  "files": [
+    "dist"
+  ],
   "scripts": {
     "dev": "modern dev",
+    "build": "tsup",
     "pretest": "modern build",
     "test": "playwright test",
     "lint:type": "tsc --noEmit"
   },
   "devDependencies": {
     "@modern-js/app-tools": "^2.46.1",
-    "@modern-js/utils": "^2.46.1",
     "@playwright/test": "^1.41.0",
+    "@types/chrome": "^0.0.258",
     "@types/node": "^20.11.5",
-    "@webx-kit/chrome-types": "workspace:^",
+    "@types/react": "^18.2.48",
+    "@types/react-dom": "^18.2.18",
     "@webx-kit/modernjs-plugin": "workspace:^",
     "@webx-kit/test-utils": "workspace:^",
+    "jotai": "^2.6.2",
+    "react": "^18.2.0",
+    "react-dom": "^18.2.0",
+    "tsup": "^8.0.1",
     "typescript": "^5.3.3"
+  },
+  "peerDependencies": {
+    "@types/chrome": "^0.0.258"
   }
 }

packages/storage/src/index.ts

@@ -1,8 +1,9 @@
 export interface ChromeStorage<T extends Record<string, any> = Record<string, any>> {
   setItem: <K extends keyof T>(key: K, value: T[K]) => Promise<void>;
-  getItem: <K extends keyof T>(key: K) => Promise<T[K] | null>;
+  getItem: <K extends keyof T>(key: K, defaultValue?: T[K]) => Promise<T[K] | null>;
   removeItem: (key: keyof T | (keyof T)[]) => Promise<void>;
   clear: () => Promise<void>;
+  subscribe: <K extends keyof T>(key: K, callback: (value: T[K] | null) => void) => VoidFunction;
 }
 
 export function createStorage<T extends Record<string, any> = Record<string, any>>(
@@ -14,9 +15,9 @@ export function createStorage<T extends Record<string, any> = Record<string, any
     return storage.set({ [key]: value });
   }
 
-  async function getItem<K extends keyof T>(key: K): Promise<T[K] | null> {
+  async function getItem<K extends keyof T>(key: K, defaultValue?: T[K]): Promise<T[K] | null> {
     const result = await storage.get(key as string);
-    return result[key as string] ?? null;
+    return result[key as string] ?? (arguments.length === 2 ? defaultValue : null);
   }
 
   async function removeItem(key: keyof T | (keyof T)[]): Promise<void> {
@@ -27,10 +28,49 @@ export function createStorage<T extends Record<string, any> = Record<string, any
     return storage.clear();
   }
 
+  let onChangedActive = false;
+  const onChangedCallbacks = new Map<keyof T, Set<(value: any) => void>>();
+  const onChangedListener: Parameters<(typeof storage)['onChanged']['addListener']>[0] = (changes) => {
+    for (const [key, value] of Object.entries(changes)) {
+      const changedCallback = onChangedCallbacks.get(key);
+      if (!changedCallback) continue;
+      const newValue = 'newValue' in value ? value.newValue : null;
+      changedCallback.forEach((callback) => callback(newValue));
+    }
+  };
+
+  function addCallback(key: string, callback: (value: any) => void) {
+    if (onChangedCallbacks.has(key)) onChangedCallbacks.get(key)!.add(callback);
+    else onChangedCallbacks.set(key, new Set([callback]));
+    if (!onChangedActive) {
+      storage.onChanged.addListener(onChangedListener);
+      onChangedActive = true;
+    }
+  }
+
+  function removeCallback(key: string, callback: (value: any) => void) {
+    const callbackSet = onChangedCallbacks.get(key);
+    if (!callbackSet) return;
+    callbackSet.delete(callback);
+    if (!callbackSet.size) onChangedCallbacks.delete(key);
+    if (!onChangedCallbacks.size) {
+      storage.onChanged.removeListener(onChangedListener);
+      onChangedActive = false;
+    }
+  }
+
+  function subscribe<K extends keyof T>(key: K, callback: (value: T[K] | null) => void): VoidFunction {
+    addCallback(key as string, callback);
+    return () => {
+      removeCallback(key as string, callback);
+    };
+  }
+
   return {
     setItem,
     getItem,
     removeItem,
     clear,
+    subscribe,
   };
 }

packages/storage/tsconfig.json

@@ -1,6 +1,6 @@
 {
   "compilerOptions": {
-    "target": "ES2015",
+    "target": "ES2020",
     "lib": ["DOM", "ESNext"],
     "allowJs": true,
     "module": "commonjs",

packages/storage/tsup.config.ts

@@ -0,0 +1,9 @@
+import { defineConfig } from 'tsup';
+
+export default defineConfig({
+  entry: ['./src/index.ts'],
+  outDir: './dist',
+  format: 'esm',
+  clean: true,
+  dts: true,
+});