[site] Preserve theme across navigation

Author: e111077

packages/lit-dev-content/src/global/theme-worker.ts

@@ -8,11 +8,12 @@ type ColorMode = 'light' | 'dark';
 
 type ThemeWorkerMessage =
   | {
-      type: 'connect' | 'system-mode';
+      type: 'connect';
       mode: ColorMode;
+      overridden: boolean;
     }
   | {
-      type: 'set-mode';
+      type: 'system-mode' | 'set-mode';
       mode: ColorMode;
     };
 
@@ -22,6 +23,7 @@ interface SharedWorkerScope {
 
 const ports = new Set<MessagePort>();
 let mode: ColorMode | undefined;
+let overridden = false;
 
 const isColorMode = (value: unknown): value is ColorMode =>
   value === 'light' || value === 'dark';
@@ -33,7 +35,7 @@ const publishMode = () => {
 
   for (const port of ports) {
     try {
-      port.postMessage({type: 'mode', mode});
+      port.postMessage({type: 'mode', mode, overridden});
     } catch {
       ports.delete(port);
     }
@@ -57,13 +59,16 @@ const publishMode = () => {
         case 'connect':
           if (mode === undefined) {
             mode = data.mode;
+            overridden = data.overridden === true;
           }
           break;
         case 'system-mode':
           mode = data.mode;
+          overridden = false;
           break;
         case 'set-mode':
           mode = data.mode;
+          overridden = true;
           break;
       }
 

packages/lit-dev-content/src/global/theme.ts

@@ -12,6 +12,7 @@ export const COLOR_MODE_CHANGE_EVENT = 'litdev-color-mode-change';
 
 interface ThemeState {
   mode: ColorMode;
+  overridden: boolean;
   mediaQuery: MediaQueryList;
   worker?: SharedWorker;
   port?: MessagePort;
@@ -20,10 +21,12 @@ interface ThemeState {
 interface ThemeWorkerModeMessage {
   type: 'mode';
   mode: ColorMode;
+  overridden: boolean;
 }
 
 const THEME_STATE_KEY = Symbol.for('lit.dev.theme-state');
 const THEME_WORKER_PATH = '/js/global/theme-worker.js';
+const COLOR_MODE_STORAGE_KEY = 'color-mode';
 
 const isColorMode = (value: unknown): value is ColorMode =>
   value === 'light' || value === 'dark';
@@ -43,14 +46,41 @@ const setThemeState = (state: ThemeState) => {
   )[THEME_STATE_KEY] = state;
 };
 
+const getStoredColorMode = (): ColorMode | undefined => {
+  try {
+    const mode = sessionStorage.getItem(COLOR_MODE_STORAGE_KEY);
+    return isColorMode(mode) ? mode : undefined;
+  } catch {
+    return undefined;
+  }
+};
+
+const storeColorMode = (mode: ColorMode | undefined) => {
+  try {
+    if (mode) {
+      sessionStorage.setItem(COLOR_MODE_STORAGE_KEY, mode);
+    } else {
+      sessionStorage.removeItem(COLOR_MODE_STORAGE_KEY);
+    }
+  } catch {
+    // Theme switching still works within this document when storage is blocked.
+  }
+};
+
 /**
  * Sets the theme on the page given a color mode.
  *
  * @param mode The source color to generate the theme.
  */
-function applyColorMode(state: ThemeState, mode: ColorMode) {
+function applyColorMode(
+  state: ThemeState,
+  mode: ColorMode,
+  overridden: boolean
+) {
   const modeChanged = state.mode !== mode;
   state.mode = mode;
+  state.overridden = overridden;
+  storeColorMode(overridden ? mode : undefined);
   document.body.classList.remove('light', 'dark', 'auto');
   document.body.classList.add(mode);
   updateMetaColor(mode);
@@ -118,20 +148,19 @@ export function initializeTheme() {
   }
 
   const mediaQuery = window.matchMedia('(prefers-color-scheme: dark)');
+  const storedMode = getStoredColorMode();
   const state: ThemeState = {
-    mode: mediaQuery.matches ? 'dark' : 'light',
+    mode: storedMode ?? (mediaQuery.matches ? 'dark' : 'light'),
+    overridden: storedMode !== undefined,
     mediaQuery,
   };
   setThemeState(state);
-  applyColorMode(state, state.mode);
+  applyColorMode(state, state.mode, state.overridden);
 
   mediaQuery.addEventListener('change', ({matches}) => {
     const mode = matches ? 'dark' : 'light';
-    if (state.port) {
-      state.port.postMessage({type: 'system-mode', mode});
-    } else {
-      applyColorMode(state, mode);
-    }
+    applyColorMode(state, mode, false);
+    state.port?.postMessage({type: 'system-mode', mode});
   });
 
   if (!('SharedWorker' in window)) {
@@ -160,10 +189,14 @@ export function initializeTheme() {
         return;
       }
 
-      applyColorMode(state, message.mode);
+      applyColorMode(state, message.mode, message.overridden === true);
     });
     port.start();
-    port.postMessage({type: 'connect', mode: state.mode});
+    port.postMessage({
+      type: 'connect',
+      mode: state.mode,
+      overridden: state.overridden,
+    });
   } catch {
     state.worker = undefined;
     state.port = undefined;
@@ -180,6 +213,6 @@ export function setColorMode(mode: ColorMode) {
     return;
   }
 
-  applyColorMode(state, mode);
+  applyColorMode(state, mode, true);
   state.port?.postMessage({type: 'set-mode', mode});
 }

packages/lit-dev-tests/src/playwright/home.spec.ts

@@ -102,6 +102,16 @@ test.describe('Home page', () => {
     await expect(firstBody).toHaveClass(/\blight\b/);
     await expect(firstToggle).toHaveAttribute('aria-pressed', 'false');
     await expect(firstToggle).toHaveAttribute('aria-label', 'Dark mode');
+    expect(
+      await page.evaluate(() => sessionStorage.getItem('color-mode'))
+    ).toBe('light');
+
+    await page.goto('/docs/');
+    await expect(firstBody).toHaveClass(/\blight\b/);
+    await expect(firstToggle).toHaveAttribute('aria-pressed', 'false');
+    expect(
+      await page.evaluate(() => sessionStorage.getItem('color-mode'))
+    ).toBe('light');
 
     const secondPage = await context.newPage();
     await secondPage.emulateMedia({colorScheme: 'dark'});
@@ -150,6 +160,9 @@ test.describe('Home page', () => {
     await waitForMediaChangeCount(2);
     await expect(firstBody).toHaveClass(/\bdark\b/);
     await expect(secondBody).toHaveClass(/\bdark\b/);
+    expect(
+      await secondPage.evaluate(() => sessionStorage.getItem('color-mode'))
+    ).toBe(null);
 
     await secondToggle.click();
     await expect(firstBody).toHaveClass(/\blight\b/);
@@ -162,6 +175,9 @@ test.describe('Home page', () => {
     await resetPage.emulateMedia({colorScheme: 'dark'});
     await resetPage.goto('/');
     await expect(resetPage.locator('body')).toHaveClass(/\bdark\b/);
+    expect(
+      await resetPage.evaluate(() => sessionStorage.getItem('color-mode'))
+    ).toBe(null);
     await expect(
       resetPage.locator(
         '#desktopNav theme-switcher button[aria-label="Dark mode"]'