refactor: replace chrome storage api with wxt's storage api (#566)

Co-authored-by: antonvishal <antonvishal@appexert.com>

apps/browser-extension/entrypoints/content/chatgpt.ts

@@ -3,9 +3,12 @@ import {
 	ELEMENT_IDS,
 	MESSAGE_TYPES,
 	POSTHOG_EVENT_KEY,
-	STORAGE_KEYS,
 	UI_CONFIG,
 } from "../../utils/constants"
+import {
+	autoSearchEnabled,
+	autoCapturePromptsEnabled,
+} from "../../utils/storage"
 import {
 	createChatGPTInputBarElement,
 	DOMUtils,
@@ -578,12 +581,9 @@ function addSaveChatGPTElementBeforeComposerBtn() {
 }
 
 async function setupChatGPTAutoFetch() {
-	const result = await chrome.storage.local.get([
-		STORAGE_KEYS.AUTO_SEARCH_ENABLED,
-	])
-	const autoSearchEnabled = result[STORAGE_KEYS.AUTO_SEARCH_ENABLED] ?? false
+	const autoSearch = (await autoSearchEnabled.getValue()) ?? false
 
-	if (!autoSearchEnabled) {
+	if (!autoSearch) {
 		return
 	}
 
@@ -640,13 +640,9 @@ function setupChatGPTPromptCapture() {
 	document.body.setAttribute("data-chatgpt-prompt-capture-setup", "true")
 
 	const capturePromptContent = async (source: string) => {
-		const result = await chrome.storage.local.get([
-			STORAGE_KEYS.AUTO_CAPTURE_PROMPTS_ENABLED,
-		])
-		const autoCapturePromptsEnabled =
-			result[STORAGE_KEYS.AUTO_CAPTURE_PROMPTS_ENABLED] ?? false
+		const autoCapture = (await autoCapturePromptsEnabled.getValue()) ?? false
 
-		if (!autoCapturePromptsEnabled) {
+		if (!autoCapture) {
 			console.log("Auto capture prompts is disabled, skipping prompt capture")
 			return
 		}

apps/browser-extension/entrypoints/content/claude.ts

@@ -3,9 +3,12 @@ import {
 	ELEMENT_IDS,
 	MESSAGE_TYPES,
 	POSTHOG_EVENT_KEY,
-	STORAGE_KEYS,
 	UI_CONFIG,
 } from "../../utils/constants"
+import {
+	autoSearchEnabled,
+	autoCapturePromptsEnabled,
+} from "../../utils/storage"
 import {
 	createClaudeInputBarElement,
 	DOMUtils,
@@ -488,13 +491,9 @@ function setupClaudePromptCapture() {
 	}
 	document.body.setAttribute("data-claude-prompt-capture-setup", "true")
 	const captureClaudePromptContent = async (source: string) => {
-		const result = await chrome.storage.local.get([
-			STORAGE_KEYS.AUTO_CAPTURE_PROMPTS_ENABLED,
-		])
-		const autoCapturePromptsEnabled =
-			result[STORAGE_KEYS.AUTO_CAPTURE_PROMPTS_ENABLED] ?? false
+		const autoCapture = (await autoCapturePromptsEnabled.getValue()) ?? false
 
-		if (!autoCapturePromptsEnabled) {
+		if (!autoCapture) {
 			console.log("Auto capture prompts is disabled, skipping prompt capture")
 			return
 		}
@@ -600,11 +599,8 @@ function setupClaudePromptCapture() {
 }
 
 async function setupClaudeAutoFetch() {
-	const result = await chrome.storage.local.get([
-		STORAGE_KEYS.AUTO_SEARCH_ENABLED,
-	])
-	const autoSearchEnabled = result[STORAGE_KEYS.AUTO_SEARCH_ENABLED] ?? false
-	if (!autoSearchEnabled) {
+	const autoSearch = (await autoSearchEnabled.getValue()) ?? false
+	if (!autoSearch) {
 		return
 	}
 

apps/browser-extension/entrypoints/content/shared.ts

@@ -1,4 +1,5 @@
-import { MESSAGE_TYPES, STORAGE_KEYS } from "../../utils/constants"
+import { MESSAGE_TYPES } from "../../utils/constants"
+import { bearerToken, userData } from "../../utils/storage"
 import { DOMUtils } from "../../utils/ui-components"
 import { default as TurndownService } from "turndown"
 
@@ -95,13 +96,13 @@ export function setupGlobalKeyboardShortcut() {
 }
 
 export function setupStorageListener() {
-	window.addEventListener("message", (event) => {
+	window.addEventListener("message", async (event) => {
 		if (event.source !== window) {
 			return
 		}
-		const bearerToken = event.data.token
-		const userData = event.data.userData
-		if (bearerToken && userData) {
+		const token = event.data.token
+		const user = event.data.userData
+		if (token && user) {
 			if (
 				!(
 					window.location.hostname === "localhost" ||
@@ -115,13 +116,11 @@ export function setupStorageListener() {
 				return
 			}
 
-			chrome.storage.local.set(
-				{
-					[STORAGE_KEYS.BEARER_TOKEN]: bearerToken,
-					[STORAGE_KEYS.USER_DATA]: userData,
-				},
-				() => {},
-			)
+			try {
+				await Promise.all([bearerToken.setValue(token), userData.setValue(user)])
+			} catch {
+				// Do nothing
+			}
 		}
 	})
 }

apps/browser-extension/entrypoints/content/t3.ts

@@ -3,9 +3,12 @@ import {
 	ELEMENT_IDS,
 	MESSAGE_TYPES,
 	POSTHOG_EVENT_KEY,
-	STORAGE_KEYS,
 	UI_CONFIG,
 } from "../../utils/constants"
+import {
+	autoSearchEnabled,
+	autoCapturePromptsEnabled,
+} from "../../utils/storage"
 import { createT3InputBarElement, DOMUtils } from "../../utils/ui-components"
 
 let t3DebounceTimeout: NodeJS.Timeout | null = null
@@ -497,13 +500,9 @@ function setupT3PromptCapture() {
 	document.body.setAttribute("data-t3-prompt-capture-setup", "true")
 
 	const captureT3PromptContent = async (source: string) => {
-		const result = await chrome.storage.local.get([
-			STORAGE_KEYS.AUTO_CAPTURE_PROMPTS_ENABLED,
-		])
-		const autoCapturePromptsEnabled =
-			result[STORAGE_KEYS.AUTO_CAPTURE_PROMPTS_ENABLED] ?? false
+		const autoCapture = (await autoCapturePromptsEnabled.getValue()) ?? false
 
-		if (!autoCapturePromptsEnabled) {
+		if (!autoCapture) {
 			console.log("Auto capture prompts is disabled, skipping prompt capture")
 			return
 		}
@@ -677,12 +676,9 @@ function setupT3PromptCapture() {
 }
 
 async function setupT3AutoFetch() {
-	const result = await chrome.storage.local.get([
-		STORAGE_KEYS.AUTO_SEARCH_ENABLED,
-	])
-	const autoSearchEnabled = result[STORAGE_KEYS.AUTO_SEARCH_ENABLED] ?? false
+	const autoSearch = (await autoSearchEnabled.getValue()) ?? false
 
-	if (!autoSearchEnabled) {
+	if (!autoSearch) {
 		return
 	}
 

apps/browser-extension/entrypoints/popup/App.tsx

@@ -2,13 +2,20 @@ import { useQueryClient } from "@tanstack/react-query"
 import { useEffect, useState } from "react"
 import "./App.css"
 import { validateAuthToken } from "../../utils/api"
-import { MESSAGE_TYPES, STORAGE_KEYS } from "../../utils/constants"
+import { MESSAGE_TYPES } from "../../utils/constants"
 import {
 	useDefaultProject,
 	useProjects,
 	useSetDefaultProject,
 	useUserData,
 } from "../../utils/query-hooks"
+import {
+	autoSearchEnabled as autoSearchEnabledStorage,
+	autoCapturePromptsEnabled as autoCapturePromptsEnabledStorage,
+	bearerToken,
+	defaultProject as defaultProjectStorage,
+	userData as userDataStorage,
+} from "../../utils/storage"
 import type { Project } from "../../utils/types"
 
 const Tooltip = ({
@@ -93,12 +100,13 @@ function App() {
 	useEffect(() => {
 		const checkAuthStatus = async () => {
 			try {
-				const result = await chrome.storage.local.get([
-					STORAGE_KEYS.BEARER_TOKEN,
-					STORAGE_KEYS.AUTO_SEARCH_ENABLED,
-					STORAGE_KEYS.AUTO_CAPTURE_PROMPTS_ENABLED,
+				const [token, autoSearch, autoCapturePrompts] = await Promise.all([
+					bearerToken.getValue(),
+					autoSearchEnabledStorage.getValue(),
+					autoCapturePromptsEnabledStorage.getValue(),
 				])
-				const hasToken = !!result[STORAGE_KEYS.BEARER_TOKEN]
+
+				const hasToken = !!token
 
 				if (hasToken) {
 					const isTokenValid = await validateAuthToken()
@@ -107,10 +115,10 @@ function App() {
 						setUserSignedIn(true)
 						setAuthInvalidated(false)
 					} else {
-						await chrome.storage.local.remove([
-							STORAGE_KEYS.BEARER_TOKEN,
-							STORAGE_KEYS.USER_DATA,
-							STORAGE_KEYS.DEFAULT_PROJECT,
+						await Promise.all([
+							bearerToken.removeValue(),
+							userDataStorage.removeValue(),
+							defaultProjectStorage.removeValue(),
 						])
 						queryClient.clear()
 						setUserSignedIn(false)
@@ -121,13 +129,8 @@ function App() {
 					setAuthInvalidated(false)
 				}
 
-				const autoSearchSetting =
-					result[STORAGE_KEYS.AUTO_SEARCH_ENABLED] ?? false
-				setAutoSearchEnabled(autoSearchSetting)
-
-				const autoCapturePromptsSetting =
-					result[STORAGE_KEYS.AUTO_CAPTURE_PROMPTS_ENABLED] ?? false
-				setAutoCapturePromptsEnabled(autoCapturePromptsSetting)
+				setAutoSearchEnabled(autoSearch ?? false)
+				setAutoCapturePromptsEnabled(autoCapturePrompts ?? false)
 			} catch (error) {
 				console.error("Error checking auth status:", error)
 				setUserSignedIn(false)
@@ -227,9 +230,7 @@ function App() {
 
 	const handleAutoSearchToggle = async (enabled: boolean) => {
 		try {
-			await chrome.storage.local.set({
-				[STORAGE_KEYS.AUTO_SEARCH_ENABLED]: enabled,
-			})
+			await autoSearchEnabledStorage.setValue(enabled)
 			setAutoSearchEnabled(enabled)
 		} catch (error) {
 			console.error("Error updating auto search setting:", error)
@@ -238,9 +239,7 @@ function App() {
 
 	const handleAutoCapturePromptsToggle = async (enabled: boolean) => {
 		try {
-			await chrome.storage.local.set({
-				[STORAGE_KEYS.AUTO_CAPTURE_PROMPTS_ENABLED]: enabled,
-			})
+			await autoCapturePromptsEnabledStorage.setValue(enabled)
 			setAutoCapturePromptsEnabled(enabled)
 		} catch (error) {
 			console.error("Error updating auto capture prompts setting:", error)
@@ -249,9 +248,11 @@ function App() {
 
 	const handleSignOut = async () => {
 		try {
-			await chrome.storage.local.remove([STORAGE_KEYS.BEARER_TOKEN])
-			await chrome.storage.local.remove([STORAGE_KEYS.USER_DATA])
-			await chrome.storage.local.remove([STORAGE_KEYS.DEFAULT_PROJECT])
+			await Promise.all([
+				bearerToken.removeValue(),
+				userDataStorage.removeValue(),
+				defaultProjectStorage.removeValue(),
+			])
 			setUserSignedIn(false)
 			queryClient.clear()
 		} catch (error) {

apps/browser-extension/utils/api.ts

@@ -1,7 +1,8 @@
 /**
  * API service for supermemory browser extension
  */
-import { API_ENDPOINTS, STORAGE_KEYS } from "./constants"
+import { API_ENDPOINTS } from "./constants"
+import { bearerToken, defaultProject, userData } from "./storage"
 import {
 	AuthenticationError,
 	type MemoryPayload,
@@ -14,8 +15,7 @@ import {
  * Get bearer token from storage
  */
 async function getBearerToken(): Promise<string> {
-	const result = await chrome.storage.local.get([STORAGE_KEYS.BEARER_TOKEN])
-	const token = result[STORAGE_KEYS.BEARER_TOKEN]
+	const token = await bearerToken.getValue()
 
 	if (!token) {
 		throw new AuthenticationError("Bearer token not found")
@@ -75,10 +75,8 @@ export async function fetchProjects(): Promise<Project[]> {
  */
 export async function getDefaultProject(): Promise<Project | null> {
 	try {
-		const result = await chrome.storage.local.get([
-			STORAGE_KEYS.DEFAULT_PROJECT,
-		])
-		return result[STORAGE_KEYS.DEFAULT_PROJECT] || null
+		const defaultProjectValue = await defaultProject.getValue()
+		return defaultProjectValue || null
 	} catch (error) {
 		console.error("Failed to get default project:", error)
 		return null
@@ -90,9 +88,7 @@ export async function getDefaultProject(): Promise<Project | null> {
  */
 export async function setDefaultProject(project: Project): Promise<void> {
 	try {
-		await chrome.storage.local.set({
-			[STORAGE_KEYS.DEFAULT_PROJECT]: project,
-		})
+		await defaultProject.setValue(project)
 	} catch (error) {
 		console.error("Failed to set default project:", error)
 		throw error
@@ -120,8 +116,7 @@ export async function validateAuthToken(): Promise<boolean> {
  */
 export async function getUserData(): Promise<{ email?: string } | null> {
 	try {
-		const result = await chrome.storage.local.get([STORAGE_KEYS.USER_DATA])
-		return result[STORAGE_KEYS.USER_DATA] || null
+		return (await userData.getValue()) || null
 	} catch (error) {
 		console.error("Failed to get user data:", error)
 		return null

apps/browser-extension/utils/constants.ts

@@ -10,21 +10,6 @@ export const API_ENDPOINTS = {
 		: "http://localhost:3000",
 } as const
 
-/**
- * Storage Keys
- */
-export const STORAGE_KEYS = {
-	BEARER_TOKEN: "bearer-token",
-	USER_DATA: "user-data",
-	TOKENS_LOGGED: "tokens-logged",
-	TWITTER_COOKIE: "twitter-cookie",
-	TWITTER_CSRF: "twitter-csrf",
-	TWITTER_AUTH_TOKEN: "twitter-auth-token",
-	DEFAULT_PROJECT: "sm-default-project",
-	AUTO_SEARCH_ENABLED: "sm-auto-search-enabled",
-	AUTO_CAPTURE_PROMPTS_ENABLED: "sm-auto-capture-prompts-enabled",
-} as const
-
 /**
  * DOM Element IDs
  */

apps/browser-extension/utils/posthog.ts

@@ -1,15 +1,14 @@
 import { PostHog } from "posthog-js/dist/module.no-external"
-import { STORAGE_KEYS } from "./constants"
+import { userData } from "./storage"
 
 export async function identifyUser(posthog: PostHog): Promise<void> {
-	const stored = await chrome.storage.local.get([STORAGE_KEYS.USER_DATA])
-	const userData = stored[STORAGE_KEYS.USER_DATA]
+	const storedUserData = await userData.getValue()
 
-	if (userData?.userId) {
-		posthog.identify(userData.userId, {
-			email: userData.email,
-			name: userData.name,
-			userId: userData.userId,
+	if (storedUserData?.userId) {
+		posthog.identify(storedUserData.userId, {
+			email: storedUserData.email,
+			name: storedUserData.name,
+			userId: storedUserData.userId,
 		})
 	}
 }

apps/browser-extension/utils/storage.ts

@@ -0,0 +1,121 @@
+/**
+ * Centralized storage layer using WXT's built-in storage API
+ */
+
+import { storage } from '#imports';
+import type { Project } from "./types"
+
+/**
+ * User authentication and profile data
+ */
+export interface UserData {
+	userId?: string
+	email?: string
+	name?: string
+}
+
+/**
+ * Twitter authentication tokens for API requests
+ */
+export interface TwitterAuthTokens {
+	cookie: string
+	csrf: string
+	auth: string
+}
+
+/**
+ * Local Storage Items (persistent across sessions)
+ */
+export const bearerToken = storage.defineItem<string>("local:bearer-token")
+
+export const userData = storage.defineItem<UserData>("local:user-data")
+
+export const defaultProject = storage.defineItem<Project>(
+	"local:sm-default-project",
+)
+
+export const autoSearchEnabled = storage.defineItem<boolean>(
+	"local:sm-auto-search-enabled",
+	{
+		fallback: false,
+	},
+)
+
+export const autoCapturePromptsEnabled = storage.defineItem<boolean>(
+	"local:sm-auto-capture-prompts-enabled",
+	{
+		fallback: false,
+	},
+)
+
+/**
+ * Session Storage Items (cleared when browser closes)
+ */
+export const tokensLogged = storage.defineItem<boolean>(
+	"session:tokens-logged",
+	{
+		fallback: false,
+	},
+)
+
+export const twitterCookie = storage.defineItem<string>(
+	"session:twitter-cookie",
+)
+
+export const twitterCsrf = storage.defineItem<string>("session:twitter-csrf")
+
+export const twitterAuthToken = storage.defineItem<string>(
+	"session:twitter-auth-token",
+)
+
+/**
+ * Helper function to get Twitter authentication tokens
+ * @returns Promise resolving to tokens or null if not available
+ */
+export async function getTwitterTokens(): Promise<TwitterAuthTokens | null> {
+	const [cookie, csrf, auth] = await Promise.all([
+		twitterCookie.getValue(),
+		twitterCsrf.getValue(),
+		twitterAuthToken.getValue(),
+	])
+
+	if (!cookie || !csrf || !auth) {
+		return null
+	}
+
+	return {
+		cookie,
+		csrf,
+		auth,
+	}
+}
+
+/**
+ * Helper function to set Twitter authentication tokens
+ * @param tokens - Twitter authentication tokens to store
+ */
+export async function setTwitterTokens(
+	tokens: TwitterAuthTokens,
+): Promise<void> {
+	await Promise.all([
+		twitterCookie.setValue(tokens.cookie),
+		twitterCsrf.setValue(tokens.csrf),
+		twitterAuthToken.setValue(tokens.auth),
+	])
+}
+
+/**
+ * Helper function to check if tokens have been logged (for one-time logging)
+ * @returns Promise resolving to boolean indicating if tokens were previously logged
+ */
+export async function getTokensLogged(): Promise<boolean> {
+	return (await tokensLogged.getValue()) ?? false
+}
+
+/**
+ * Helper function to mark tokens as logged
+ */
+export async function setTokensLogged(): Promise<void> {
+	await tokensLogged.setValue(true)
+}
+

apps/browser-extension/utils/twitter-auth.ts

@@ -2,24 +2,23 @@
  * Twitter Authentication Module
  * Handles token capture and storage for Twitter API access
  */
-import { STORAGE_KEYS } from "./constants"
-
-export interface TwitterAuthTokens {
-	cookie: string
-	csrf: string
-	auth: string
-}
+import {
+	getTokensLogged,
+	setTokensLogged,
+	setTwitterTokens,
+	type TwitterAuthTokens,
+} from "./storage"
 
 /**
  * Captures Twitter authentication tokens from web request headers
  * @param details - Web request details containing headers
  * @returns True if tokens were captured, false otherwise
  */
-export function captureTwitterTokens(
+export async function captureTwitterTokens(
 	details: chrome.webRequest.WebRequestDetails & {
 		requestHeaders?: chrome.webRequest.HttpHeader[]
 	},
-): boolean {
+): Promise<boolean> {
 	if (!(details.url.includes("x.com") || details.url.includes("twitter.com"))) {
 		return false
 	}
@@ -35,17 +34,16 @@ export function captureTwitterTokens(
 	)
 
 	if (authHeader?.value && cookieHeader?.value && csrfHeader?.value) {
-		chrome.storage.session.get([STORAGE_KEYS.TOKENS_LOGGED], (result) => {
-			if (!result[STORAGE_KEYS.TOKENS_LOGGED]) {
-				console.log("Twitter auth tokens captured successfully")
-				chrome.storage.session.set({ [STORAGE_KEYS.TOKENS_LOGGED]: true })
-			}
-		})
+		const tokensAlreadyLogged = await getTokensLogged()
+		if (!tokensAlreadyLogged) {
+			console.log("Twitter auth tokens captured successfully")
+			await setTokensLogged()
+		}
 
-		chrome.storage.session.set({
-			[STORAGE_KEYS.TWITTER_COOKIE]: cookieHeader.value,
-			[STORAGE_KEYS.TWITTER_CSRF]: csrfHeader.value,
-			[STORAGE_KEYS.TWITTER_AUTH_TOKEN]: authHeader.value,
+		await setTwitterTokens({
+			cookie: cookieHeader.value,
+			csrf: csrfHeader.value,
+			auth: authHeader.value,
 		})
 
 		return true
@@ -54,32 +52,6 @@ export function captureTwitterTokens(
 	return false
 }
 
-/**
- * Retrieves stored Twitter authentication tokens
- * @returns Promise resolving to tokens or null if not available
- */
-export async function getTwitterTokens(): Promise<TwitterAuthTokens | null> {
-	const result = await chrome.storage.session.get([
-		STORAGE_KEYS.TWITTER_COOKIE,
-		STORAGE_KEYS.TWITTER_CSRF,
-		STORAGE_KEYS.TWITTER_AUTH_TOKEN,
-	])
-
-	if (
-		!result[STORAGE_KEYS.TWITTER_COOKIE] ||
-		!result[STORAGE_KEYS.TWITTER_CSRF] ||
-		!result[STORAGE_KEYS.TWITTER_AUTH_TOKEN]
-	) {
-		return null
-	}
-
-	return {
-		cookie: result[STORAGE_KEYS.TWITTER_COOKIE],
-		csrf: result[STORAGE_KEYS.TWITTER_CSRF],
-		auth: result[STORAGE_KEYS.TWITTER_AUTH_TOKEN],
-	}
-}
-
 /**
  * Creates HTTP headers for Twitter API requests using stored tokens
  * @param tokens - Twitter authentication tokens

apps/browser-extension/utils/twitter-import.ts

@@ -4,7 +4,8 @@
  */
 import { saveAllTweets } from "./api"
 import type { MemoryPayload } from "./types"
-import { createTwitterAPIHeaders, getTwitterTokens } from "./twitter-auth"
+import { createTwitterAPIHeaders } from "./twitter-auth"
+import { getTwitterTokens } from "./storage"
 import {
 	BOOKMARKS_URL,
 	BOOKMARK_COLLECTION_URL,