Building a Chrome Extension Using React and Vite: Part 2 - State Management and Message Passing

Welcome back to our ongoing series on building a Chrome extension with React and Vite. In the first part of this tutorial series, we got started with our Chrome extension setup. If you've missed that, you can check it out here. This time, we will dive into the heart of our extension, state management, and understand how our extension's components interact with each other.

Before we move forward, let's pause and clarify the purpose of state management in our extension. State management refers to how we store and manipulate a data source accessible across various parts of our application. It's a crucial aspect of most applications as it helps us control data flow and communication throughout our application.

Our extension consists of three parts: the popup, the content script, and the background script.

• The popup is our control panel for the extension, allowing us to toggle its features on and off.
• The content script is responsible for modifying the website's DOM by adding checkboxes, buttons, and other UI elements as necessary.
• The background script, although not mentioned in the first part, serves as a centralized data management hub, acting as our single source of truth.

You may ask, "Why do we need a background script? Can't we directly connect the popup and content script to local storage changes?" Technically, you could. However, there are a few reasons why this isn't the best practice:

1. Performance: Directly connecting multiple parts of your extension to the local storage can lead to performance issues. With a background script, we can centralize state management and efficiently communicate changes.
2. Complexity: Trying to manage and track changes across various parts of your extension can quickly get complicated. A background script simplifies state management and makes debugging easier.
3. Best Practice: Utilizing a background script for managing state in your extension is a common pattern in extension development. It adheres to the principles of unidirectional data flow, a core concept in React development. This pattern ensures data consistency and simplifies data flow throughout your application.
4. Persistence: The background script maintains its state as long as the browser is running, ensuring the extension state's persistence. It manages and responds to events from all parts of your extension consistently.

Understanding these points, it's clear that our background script plays a crucial role in coordinating state and communication between the popup and the content script. By treating the background script as our state management layer, we can keep our UI - the popup and the content script - focused on presentation and user interaction.

Let's delve into the technical aspect of our blog post now.

So, how do we establish communication between these parts? It's pretty straightforward. The background script needs to store the state. We can use various data structures to do this, such as an object, a map, a set, or others depending on your specific needs. For this tutorial, I decided to use Valtio, a simple state management library - not for any special reason, but simply to experiment with it.

The background script needs to send the state to the requested party, as well as update the state upon request. That sounds like sending and receiving messages, doesn't it?

Firstly, let's handle the state. Here are the steps:

1. Create a file extensionState.ts
2. Create a state inside the file:

import { proxy } from "valtio";

export enum Actions {
  GET_STATE = "get-state",
  SET_STATE = "set-state",
}
export type Action = `${Actions}`;

export const state = proxy({
  isBulkDeleteEnabled: false,
  isChatExportEnabled: false,
  isPromptSavingEnabled: false,
});

export type ExtensionState = typeof state;

export const updateState = (payload: Partial<typeof state>) => {
  chrome.runtime.sendMessage({ type: Actions.SET_STATE, payload });
};

I also added a few types and an action to update the state.

Next, onto background script. First, when the extension is first installed, we need to initialize our state in the local storage:

chrome.runtime.onInstalled.addListener(() => {
  chrome.storage.local.set(state);
});

Good. Now we can persist the state between page reloads, and we can also access the state from other contexts as well.

Then, we need to send the state to those who ask for it, as well as implement changes to the state:

import { Actions, ExtensionState, state } from "../state/extensionState";

interface MessageWithoutPayload {
  type: Actions.GET_STATE;
  payload?: never;
}

interface MessageWithPayload {
  type: Actions.SET_STATE;
  payload: Partial<ExtensionState>;
}

type Message = MessageWithoutPayload | MessageWithPayload;

chrome.runtime.onInstalled.addListener(() => {
  chrome.storage.local.set(state);

  chrome.runtime.onMessage.addListener((message: Message, _, sendResponse) => {
    if (message.type === Actions.GET_STATE) {
      sendResponse(state);
    }

    if (message.type === Actions.SET_STATE) {
      Object.assign(state, message.payload);
      chrome.storage.local.set(state);
    }
  });
});

When we receive the GET_STATE message, we respond with the local valtio state because it's in sync with the local storage.
When we receive the SET_STATE message, we mutate our local state and update the local storage.

Background worker is now ready to share and update the state. How cool is that?

How about consuming the state? For this we will create a useExtensionState hook that will return the extension state. First, we get a snapshot of the state. Snapshot is a valtio's way to provide reactive state of our state (sorry for the tautology 😅):

import { useSnapshot } from "valtio";

const useExtensionState = () => {
  const extensionState = useSnapshot(state);

  return extensionState;
};

Cool. Now we have access to the extension state. But it is not the same state that exists in the background context. Or in the popup context. Or in the content script context. They all live their own independent life! So we need to sync the state. Let's use useEffect to handle side effects of retrieving the state from the local storage, or rather asking the background script to get the state for us:

import { useEffect } from "react";
import { Actions, state } from "../state/extensionState";
import { useSnapshot } from "valtio";

const useExtensionState = () => {
  const extensionState = useSnapshot(state);

  useEffect(() => {
    chrome.runtime.sendMessage({ type: Actions.GET_STATE }, (response) => {
      Object.assign(state, response);
    });
  }, []);

  return extensionState;
};

Here, we're asking the background script to get us the current state of our application, and update the local state with the retrieved value. And now we have access to the actual state of our extension. I hear someone yell "Hoozah!" — not so fast.

When the state changes, it won't be reflected in our local state. We can only receive the updated state whenever we refresh the page, which is not good, because we want REACTIVITY.

I tried subscribing to the state changes using valtio's subscribe, and whenever the state changes send notification to all parties that the state is updated, but the content script can't listen to messages sent via chrome.runtime.sendMessage, so I had two options:

1. Make the background script a bit more complicated, finding the active tab and sending message to that tab
2. Subscribe to local storage changes inside the useExtensionState hook

The second option looks easier to me so I went with it:

import { useEffect } from "react";
import { Actions, ExtensionState, state } from "../state/extensionState";
import { useSnapshot } from "valtio";

const useExtensionState = () => {
  const extensionState = useSnapshot(state);

  useEffect(() => {
    chrome.runtime.sendMessage({ type: Actions.GET_STATE }, (response) => {
      Object.assign(state, response);
    });
    const listener = (changes: {
      [key: string]: chrome.storage.StorageChange;
    }) => {
      Object.assign(
        state,
        Object.entries(changes).reduce((acc, [key, { newValue }]) => {
          acc[key as keyof ExtensionState] = newValue;
          return acc;
        }, {} as typeof state)
      );
    };

    chrome.storage.onChanged.addListener(listener);

    return () => {
      chrome.storage.onChanged.removeListener(listener);
    };
  }, []);

  return extensionState;
};

The changes is an object whose values are also an object with oldValue and newValue. Because of that, we can't directly assign our state these changes, so we have to do a little bit of data transformation. We take key-value pair from the object, and assign each key a new state value. We merge the result into our state, and that's it. We also add a cleanup function to remove unneeded listeners.

Let's test this out!

I changed the popup's App component like this:

import useExtensionState from "../hooks/useExtensionState";
import { updateState } from "../state/extensionState";

function App() {
  const extensionState = useExtensionState();

  return (
    <>
      <button
        onClick={() => {
          updateState({
            isBulkDeleteEnabled: !extensionState.isBulkDeleteEnabled,
          });
        }}
      >
        {extensionState.isBulkDeleteEnabled ? "Disable" : "Enable"}
      </button>
    </>
  );
}

Whenever the button is clicked, it should toggle the isBulkDeleteEnabled value from true to false, and vice versa. To know that it's working inside the button we render text depending on the state value.

We do the same for content script with one exception - we don't want to toggle state from content script (though you can if you want), we just want to make sure it has up to date state on each change:

import useExtensionState from "../hooks/useExtensionState";

function App() {
  const extensionState = useExtensionState();
  return (
    <>
      <div>{extensionState.isBulkDeleteEnabled ? "ENABLED" : "DISABLED"}</div>
    </>
  );
}

Now, rebuild the extension yarn build, and click the refresh button in the extensions tab. Open the https://chat.openai.com/, toggle the extension by clicking its button and try clicking the button. In the popup, the Disable text switches with Enable, and in the content script (which is barely visible due to its black text color and transparent background) also changes its text.

Now you can yell "Hoozah!".

What's left is to develop a pure react application for both popup and the content script, and so I decided to conclude the tutorial on the extension development with React.

If you want to follow along to the bitter (hopefully sweeter) end, please let me know in the comments and I will continue the series.

Hope you enjoyed this series. See you next time 😊