From 8bcd7aa3bad4a21a64e3f4febda59993cd4749dc Mon Sep 17 00:00:00 2001
From: Kevin Smith <zenparsing@gmail.com>
Date: Tue, 16 Dec 2025 07:55:18 -0500
Subject: [PATCH] [NTP] Simplify state provider API (#32643)

---
 .../lib/state_provider.test.tsx               | 158 ++++++++++++++++++
 .../lib/state_provider.ts                     |  99 ++++++++++-
 .../lib/{store.ts => state_store.ts}          |   8 +-
 .../state/background_handler.ts               |   4 +-
 .../state/new_tab_handler.ts                  |   6 +-
 .../state/rewards_handler.ts                  |   4 +-
 .../state/search_handler.ts                   |   6 +-
 .../state/top_sites_handler.ts                |   4 +-
 .../state/vpn_handler.ts                      |   4 +-
 .../stories/background_handler.ts             |   4 +-
 .../stories/new_tab_handler.ts                |   6 +-
 .../stories/rewards_handler.ts                |   4 +-
 .../stories/search_handler.ts                 |   6 +-
 .../stories/top_sites_handler.ts              |   4 +-
 .../stories/vpn_handler.ts                    |   4 +-
 15 files changed, 285 insertions(+), 36 deletions(-)
 create mode 100644 browser/resources/brave_new_tab_page_refresh/lib/state_provider.test.tsx
 rename browser/resources/brave_new_tab_page_refresh/lib/{store.ts => state_store.ts} (94%)

diff --git a/browser/resources/brave_new_tab_page_refresh/lib/state_provider.test.tsx b/browser/resources/brave_new_tab_page_refresh/lib/state_provider.test.tsx
new file mode 100644
index 00000000000..5d95eb3ab0f
--- /dev/null
+++ b/browser/resources/brave_new_tab_page_refresh/lib/state_provider.test.tsx
@@ -0,0 +1,158 @@
+/* Copyright (c) 2025 The Brave Authors. All rights reserved.
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this file,
+ * You can obtain one at https://mozilla.org/MPL/2.0/. */
+
+import * as React from 'react'
+import { render, screen, act } from '@testing-library/react'
+
+import { createStateProvider } from './state_provider'
+import { StateStore } from './state_store'
+
+interface TestState {
+  count: number
+  name: string
+}
+
+function defaultState(): TestState {
+  return { count: 0, name: 'test' }
+}
+
+interface TestActions {
+  increment(): void
+  setName(name: string): void
+}
+
+function createHandler(store: StateStore<TestState>): TestActions {
+  return {
+    increment() {
+      store.update((s) => ({ count: s.count + 1 }))
+    },
+    setName(name: string) {
+      store.update({ name })
+    },
+  }
+}
+
+function flushMicrotasks() {
+  return act(() => new Promise<void>((resolve) => queueMicrotask(resolve)))
+}
+
+describe('createStateProvider', () => {
+  describe('Provider', () => {
+    it('renders children', () => {
+      const TestProvider = createStateProvider(defaultState(), createHandler)
+      render(
+        <TestProvider>
+          <div data-testid='child'>Hello</div>
+        </TestProvider>,
+      )
+      expect(screen.getByTestId('child')).toHaveTextContent('Hello')
+    })
+
+    it('exposes store to window.appState when name prop is provided', () => {
+      const TestProvider = createStateProvider(defaultState(), createHandler)
+      render(
+        <TestProvider name='test'>
+          <div />
+        </TestProvider>,
+      )
+      expect((self as any).appState.test).toBeDefined()
+      expect((self as any).appState.test.getState()).toEqual(defaultState())
+    })
+  })
+
+  describe('useState', () => {
+    it('returns the mapped state value', () => {
+      const TestProvider = createStateProvider(defaultState(), createHandler)
+
+      function TestComponent() {
+        const count = TestProvider.useState((s) => s.count)
+        return <div data-testid='count'>{count}</div>
+      }
+
+      render(
+        <TestProvider>
+          <TestComponent />
+        </TestProvider>,
+      )
+      expect(screen.getByTestId('count')).toHaveTextContent('0')
+    })
+
+    it('updates when state changes', async () => {
+      const TestProvider = createStateProvider(defaultState(), createHandler)
+
+      function TestComponent() {
+        const count = TestProvider.useState((s) => s.count)
+        const actions = TestProvider.useActions()
+        return (
+          <button
+            data-testid='button'
+            onClick={() => actions.increment()}
+          >
+            {count}
+          </button>
+        )
+      }
+
+      render(
+        <TestProvider>
+          <TestComponent />
+        </TestProvider>,
+      )
+
+      expect(screen.getByTestId('button')).toHaveTextContent('0')
+      screen.getByTestId('button').click()
+      await flushMicrotasks()
+      expect(screen.getByTestId('button')).toHaveTextContent('1')
+    })
+
+    it('throws when used outside provider', () => {
+      const TestProvider = createStateProvider(defaultState(), createHandler)
+
+      function TestComponent() {
+        TestProvider.useState((s) => s.count)
+        return null
+      }
+
+      expect(() => render(<TestComponent />)).toThrow(
+        'State context value has not been set',
+      )
+    })
+  })
+
+  describe('useActions', () => {
+    it('returns the actions object', () => {
+      const TestProvider = createStateProvider(defaultState(), createHandler)
+      let capturedActions: TestActions | null = null
+
+      function TestComponent() {
+        capturedActions = TestProvider.useActions()
+        return null
+      }
+
+      render(
+        <TestProvider>
+          <TestComponent />
+        </TestProvider>,
+      )
+
+      expect(capturedActions).not.toBeNull()
+      expect(typeof capturedActions!.increment).toBe('function')
+      expect(typeof capturedActions!.setName).toBe('function')
+    })
+
+    it('throws when used outside provider', () => {
+      const TestProvider = createStateProvider(defaultState(), createHandler)
+
+      function TestComponent() {
+        TestProvider.useActions()
+        return null
+      }
+
+      expect(() => render(<TestComponent />)).toThrow(
+        'State context value has not been set',
+      )
+    })
+  })
+})
diff --git a/browser/resources/brave_new_tab_page_refresh/lib/state_provider.ts b/browser/resources/brave_new_tab_page_refresh/lib/state_provider.ts
index 35fb05e6daf..b2a882c2209 100644
--- a/browser/resources/brave_new_tab_page_refresh/lib/state_provider.ts
+++ b/browser/resources/brave_new_tab_page_refresh/lib/state_provider.ts
@@ -5,15 +5,96 @@
 
 import * as React from 'react'
 
-import { Store, createStore } from './store'
+import { StateStore, createStateStore } from './state_store'
 
+/**
+ * Creates a React context provider component for managing application state.
+ * The returned Provider component has `useState` and `useActions` hooks
+ * attached as static methods for accessing state and actions within the
+ * provider tree.
+ *
+ * @param initialState - The initial state object for the store.
+ * @param createHandler - A function that receives the state store and returns
+ *   an action handler. Action handlers typically call `store.update()` to
+ *   modify state.
+ * @returns A Provider component with attached `useState` and `useActions`
+ *   hooks.
+ *
+ * @example
+ * // 1. Define state and actions:
+ *
+ * interface AppState {
+ *   count: number
+ *   name: string
+ * }
+ *
+ * function defaultState(): AppState {
+ *   return { count: 0, name: 'test' }
+ * }
+ *
+ * interface AppActions {
+ *   increment(): void
+ *   setName(name: string): void
+ * }
+ *
+ * function createHandler(store: StateStore<AppState>): AppActions {
+ *   return {
+ *     increment() {
+ *       store.update((s) => ({ count: s.count + 1 }))
+ *     },
+ *     setName(name) {
+ *       store.update({ name })
+ *     },
+ *   }
+ * }
+ *
+ * // 2. Create the provider:
+ *
+ * export const AppStateProvider = createStateProvider(
+ *   defaultState(),
+ *   createHandler
+ * )
+ *
+ * // Export hooks for convenience:
+ *
+ * export const useAppState = AppStateProvider.useState
+ * export const useAppActions = AppStateProvider.useActions
+ *
+ * // Wrap your app with the provider. The optional `name` prop exposes the
+ * // store to `window.appState[name]` for debugging.
+ *
+ * function App() {
+ *   return (
+ *     <AppStateProvider name='myApp'>
+ *       <MyComponent />
+ *     </AppStateProvider>
+ *   )
+ * }
+ *
+ * // Use hooks in components to access state and actions:
+ *
+ * function MyComponent() {
+ *   // Select specific state values with a mapping function.
+ *   const count = useAppState((s) => s.count)
+ *   const name = useAppState((s) => s.name)
+ *
+ *   // Get actions to update state.
+ *   const actions = useAppActions()
+ *
+ *   return (
+ *     <button onClick={() => actions.increment()}>
+ *       {count}
+ *     </button>
+ *   )
+ * }
+ */
 export function createStateProvider<State, Actions>(
   initialState: State,
-  createHandler: (store: Store<State>) => Actions,
+  createHandler: (store: StateStore<State>) => Actions,
 ) {
   interface ContextValue {
-    store: Store<State>
-    handler: Actions
+    store: StateStore<State>
+    actions: Actions
   }
 
   const context = React.createContext<ContextValue | null>(null)
@@ -27,7 +108,7 @@ export function createStateProvider<State, Actions>(
   }
 
   function useActions(): Actions {
-    return useContextValue().handler
+    return useContextValue().actions
   }
 
   function useState<T>(map: (state: State) => T): T {
@@ -44,15 +125,15 @@ export function createStateProvider<State, Actions>(
 
   interface ProviderProps {
     name?: string
-    createHandler?: (store: Store<State>) => Actions
+    createHandler?: (store: StateStore<State>) => Actions
     children: React.ReactNode
   }
 
   function Provider(props: ProviderProps) {
     const value = React.useMemo(() => {
-      const store = createStore(initialState)
-      const handler = (props.createHandler ?? createHandler)(store)
-      return { store, handler }
+      const store = createStateStore(initialState)
+      const actions = (props.createHandler ?? createHandler)(store)
+      return { store, actions }
     }, [props.createHandler])
 
     React.useEffect(() => {
diff --git a/browser/resources/brave_new_tab_page_refresh/lib/store.ts b/browser/resources/brave_new_tab_page_refresh/lib/state_store.ts
similarity index 94%
rename from browser/resources/brave_new_tab_page_refresh/lib/store.ts
rename to browser/resources/brave_new_tab_page_refresh/lib/state_store.ts
index de75cd70cd6..295ec3498fe 100644
--- a/browser/resources/brave_new_tab_page_refresh/lib/store.ts
+++ b/browser/resources/brave_new_tab_page_refresh/lib/state_store.ts
@@ -8,7 +8,7 @@ type Listener<State> = (state: State) => void
 type UpdateFunction<State> = (state: State) => Partial<State>
 
 // A simple object-state store.
-export interface Store<State> {
+export interface StateStore<State> {
   // Returns the current state of the store.
   getState: () => State
 
@@ -19,11 +19,13 @@ export interface Store<State> {
 
   // Adds a listener that will be notified when the state store changes. The
   // listener will not be notified immediately. Returns a function that will
-  // remove the listener from store.
+  // remove the listener from the store.
   addListener: (listener: Listener<State>) => () => void
 }
 
-export function createStore<State>(initialState: State): Store<State> {
+export function createStateStore<State>(
+  initialState: State,
+): StateStore<State> {
   const listeners = new Set<Listener<State>>()
   const state = { ...initialState }
   let notificationQueued = false
diff --git a/browser/resources/brave_new_tab_page_refresh/state/background_handler.ts b/browser/resources/brave_new_tab_page_refresh/state/background_handler.ts
index c4264e1af89..9041a6a0a3d 100644
--- a/browser/resources/brave_new_tab_page_refresh/state/background_handler.ts
+++ b/browser/resources/brave_new_tab_page_refresh/state/background_handler.ts
@@ -6,13 +6,13 @@
 import { loadTimeData } from '$web-common/loadTimeData'
 import { SponsoredRichMediaAdEventHandler } from 'gen/brave/components/ntp_background_images/browser/mojom/ntp_background_images.mojom.m.js'
 import { NewTabPageProxy } from './new_tab_page_proxy'
-import { Store } from '../lib/store'
+import { StateStore } from '../lib/state_store'
 import { debounce } from '$web-common/debounce'
 import { preloadedBackgrounds } from './background_images/preloaded'
 import { BackgroundState, BackgroundActions } from './background_state'
 
 export function createBackgroundHandler(
-  store: Store<BackgroundState>,
+  store: StateStore<BackgroundState>,
 ): BackgroundActions {
   const newTabProxy = NewTabPageProxy.getInstance()
   const { handler } = newTabProxy
diff --git a/browser/resources/brave_new_tab_page_refresh/state/new_tab_handler.ts b/browser/resources/brave_new_tab_page_refresh/state/new_tab_handler.ts
index 85d5ff99b0a..9049f2791e5 100644
--- a/browser/resources/brave_new_tab_page_refresh/state/new_tab_handler.ts
+++ b/browser/resources/brave_new_tab_page_refresh/state/new_tab_handler.ts
@@ -5,11 +5,13 @@
 
 import { loadTimeData } from '$web-common/loadTimeData'
 import { NewTabPageProxy } from './new_tab_page_proxy'
-import { Store } from '../lib/store'
+import { StateStore } from '../lib/state_store'
 import { debounce } from '$web-common/debounce'
 import { NewTabState, NewTabActions } from './new_tab_state'
 
-export function createNewTabHandler(store: Store<NewTabState>): NewTabActions {
+export function createNewTabHandler(
+  store: StateStore<NewTabState>,
+): NewTabActions {
   const newTabProxy = NewTabPageProxy.getInstance()
   const { handler } = newTabProxy
 
diff --git a/browser/resources/brave_new_tab_page_refresh/state/rewards_handler.ts b/browser/resources/brave_new_tab_page_refresh/state/rewards_handler.ts
index f8b3f7ea37d..44b2f26a3f3 100644
--- a/browser/resources/brave_new_tab_page_refresh/state/rewards_handler.ts
+++ b/browser/resources/brave_new_tab_page_refresh/state/rewards_handler.ts
@@ -7,7 +7,7 @@ import { loadTimeData } from '$web-common/loadTimeData'
 import { RewardsPageProxy } from '../../../../components/brave_rewards/resources/rewards_page/webui/rewards_page_proxy'
 import { externalWalletFromExtensionData } from '../../../../components/brave_rewards/resources/shared/lib/external_wallet'
 import { NewTabPageProxy } from './new_tab_page_proxy'
-import { Store } from '../lib/store'
+import { StateStore } from '../lib/state_store'
 import { debounce } from '$web-common/debounce'
 import {
   RewardsState,
@@ -16,7 +16,7 @@ import {
 } from './rewards_state'
 
 export function createRewardsHandler(
-  store: Store<RewardsState>,
+  store: StateStore<RewardsState>,
 ): RewardsActions {
   if (!loadTimeData.getBoolean('rewardsFeatureEnabled')) {
     store.update({ initialized: true })
diff --git a/browser/resources/brave_new_tab_page_refresh/state/search_handler.ts b/browser/resources/brave_new_tab_page_refresh/state/search_handler.ts
index 69546edaf54..e3ed4e31793 100644
--- a/browser/resources/brave_new_tab_page_refresh/state/search_handler.ts
+++ b/browser/resources/brave_new_tab_page_refresh/state/search_handler.ts
@@ -6,7 +6,7 @@
 import { loadTimeData } from '$web-common/loadTimeData'
 import { SearchBoxProxy } from './search_box_proxy'
 import { NewTabPageProxy } from './new_tab_page_proxy'
-import { Store } from '../lib/store'
+import { StateStore } from '../lib/state_store'
 import { debounce } from '$web-common/debounce'
 
 import {
@@ -63,7 +63,9 @@ function storeEnabledSearchEngines(engines: Set<string>) {
   localStorage.setItem(enabledSearchEnginesStorageKey, JSON.stringify(record))
 }
 
-export function createSearchHandler(store: Store<SearchState>): SearchActions {
+export function createSearchHandler(
+  store: StateStore<SearchState>,
+): SearchActions {
   if (!loadTimeData.getBoolean('ntpSearchFeatureEnabled')) {
     return defaultSearchActions()
   }
diff --git a/browser/resources/brave_new_tab_page_refresh/state/top_sites_handler.ts b/browser/resources/brave_new_tab_page_refresh/state/top_sites_handler.ts
index f23bafa0567..d264242ed2a 100644
--- a/browser/resources/brave_new_tab_page_refresh/state/top_sites_handler.ts
+++ b/browser/resources/brave_new_tab_page_refresh/state/top_sites_handler.ts
@@ -10,11 +10,11 @@ import {
   TopSitesActions,
   TopSitesListKind,
 } from './top_sites_state'
-import { Store } from '../lib/store'
+import { StateStore } from '../lib/state_store'
 import { debounce } from '$web-common/debounce'
 
 export function createTopSitesHandler(
-  store: Store<TopSitesState>,
+  store: StateStore<TopSitesState>,
 ): TopSitesActions {
   const newTabProxy = NewTabPageProxy.getInstance()
   const { handler } = newTabProxy
diff --git a/browser/resources/brave_new_tab_page_refresh/state/vpn_handler.ts b/browser/resources/brave_new_tab_page_refresh/state/vpn_handler.ts
index f6d9a591265..ee1bd8865cc 100644
--- a/browser/resources/brave_new_tab_page_refresh/state/vpn_handler.ts
+++ b/browser/resources/brave_new_tab_page_refresh/state/vpn_handler.ts
@@ -6,7 +6,7 @@
 import { loadTimeData } from '$web-common/loadTimeData'
 import * as mojom from 'gen/brave/components/brave_vpn/common/mojom/brave_vpn.mojom.m'
 import { NewTabPageProxy } from './new_tab_page_proxy'
-import { Store } from '../lib/store'
+import { StateStore } from '../lib/state_store'
 import { debounce } from '$web-common/debounce'
 import {
   VpnState,
@@ -15,7 +15,7 @@ import {
   ConnectionState,
 } from './vpn_state'
 
-export function createVpnHandler(store: Store<VpnState>): VpnActions {
+export function createVpnHandler(store: StateStore<VpnState>): VpnActions {
   if (!loadTimeData.getBoolean('vpnFeatureEnabled')) {
     store.update({ initialized: true })
     return defaultVpnActions()
diff --git a/browser/resources/brave_new_tab_page_refresh/stories/background_handler.ts b/browser/resources/brave_new_tab_page_refresh/stories/background_handler.ts
index fdb17cb5041..f52d6b1ecc1 100644
--- a/browser/resources/brave_new_tab_page_refresh/stories/background_handler.ts
+++ b/browser/resources/brave_new_tab_page_refresh/stories/background_handler.ts
@@ -3,7 +3,7 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this file,
  * You can obtain one at https://mozilla.org/MPL/2.0/. */
 
-import { Store } from '../lib/store'
+import { StateStore } from '../lib/state_store'
 
 import {
   BackgroundState,
@@ -56,7 +56,7 @@ const sponsoredBackgrounds = {
 }
 
 export function createBackgroundHandler(
-  store: Store<BackgroundState>,
+  store: StateStore<BackgroundState>,
   args: StorybookArgs,
 ): BackgroundActions {
   store.update({
diff --git a/browser/resources/brave_new_tab_page_refresh/stories/new_tab_handler.ts b/browser/resources/brave_new_tab_page_refresh/stories/new_tab_handler.ts
index 28439fcd8f6..402039d44cb 100644
--- a/browser/resources/brave_new_tab_page_refresh/stories/new_tab_handler.ts
+++ b/browser/resources/brave_new_tab_page_refresh/stories/new_tab_handler.ts
@@ -3,7 +3,7 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this file,
  * You can obtain one at https://mozilla.org/MPL/2.0/. */
 
-import { Store } from '../lib/store'
+import { StateStore } from '../lib/state_store'
 
 import {
   NewTabState,
@@ -11,7 +11,9 @@ import {
   defaultNewTabActions,
 } from '../state/new_tab_state'
 
-export function createNewTabHandler(store: Store<NewTabState>): NewTabActions {
+export function createNewTabHandler(
+  store: StateStore<NewTabState>,
+): NewTabActions {
   store.update({
     initialized: true,
     showClock: true,
diff --git a/browser/resources/brave_new_tab_page_refresh/stories/rewards_handler.ts b/browser/resources/brave_new_tab_page_refresh/stories/rewards_handler.ts
index 936b4210f86..d62af4b72a5 100644
--- a/browser/resources/brave_new_tab_page_refresh/stories/rewards_handler.ts
+++ b/browser/resources/brave_new_tab_page_refresh/stories/rewards_handler.ts
@@ -3,7 +3,7 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this file,
  * You can obtain one at https://mozilla.org/MPL/2.0/. */
 
-import { Store } from '../lib/store'
+import { StateStore } from '../lib/state_store'
 
 import {
   RewardsState,
@@ -12,7 +12,7 @@ import {
 } from '../state/rewards_state'
 
 export function createRewardsHandler(
-  store: Store<RewardsState>,
+  store: StateStore<RewardsState>,
 ): RewardsActions {
   store.update({
     initialized: true,
diff --git a/browser/resources/brave_new_tab_page_refresh/stories/search_handler.ts b/browser/resources/brave_new_tab_page_refresh/stories/search_handler.ts
index 8eac2717831..da65e4a2601 100644
--- a/browser/resources/brave_new_tab_page_refresh/stories/search_handler.ts
+++ b/browser/resources/brave_new_tab_page_refresh/stories/search_handler.ts
@@ -3,14 +3,16 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this file,
  * You can obtain one at https://mozilla.org/MPL/2.0/. */
 
-import { Store } from '../lib/store'
+import { StateStore } from '../lib/state_store'
 import {
   SearchState,
   SearchActions,
   defaultSearchActions,
 } from '../state/search_state'
 
-export function createSearchHandler(store: Store<SearchState>): SearchActions {
+export function createSearchHandler(
+  store: StateStore<SearchState>,
+): SearchActions {
   store.update({
     initialized: true,
 
diff --git a/browser/resources/brave_new_tab_page_refresh/stories/top_sites_handler.ts b/browser/resources/brave_new_tab_page_refresh/stories/top_sites_handler.ts
index af75c2ef66f..5ff6a49da78 100644
--- a/browser/resources/brave_new_tab_page_refresh/stories/top_sites_handler.ts
+++ b/browser/resources/brave_new_tab_page_refresh/stories/top_sites_handler.ts
@@ -3,7 +3,7 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this file,
  * You can obtain one at https://mozilla.org/MPL/2.0/. */
 
-import { Store } from '../lib/store'
+import { StateStore } from '../lib/state_store'
 
 import {
   TopSitesState,
@@ -13,7 +13,7 @@ import {
 } from '../state/top_sites_state'
 
 export function createTopSitesHandler(
-  store: Store<TopSitesState>,
+  store: StateStore<TopSitesState>,
 ): TopSitesActions {
   let lastRemovedSite: TopSite | null = null
 
diff --git a/browser/resources/brave_new_tab_page_refresh/stories/vpn_handler.ts b/browser/resources/brave_new_tab_page_refresh/stories/vpn_handler.ts
index 9b319d65f8f..927e620e713 100644
--- a/browser/resources/brave_new_tab_page_refresh/stories/vpn_handler.ts
+++ b/browser/resources/brave_new_tab_page_refresh/stories/vpn_handler.ts
@@ -3,7 +3,7 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this file,
  * You can obtain one at https://mozilla.org/MPL/2.0/. */
 
-import { Store } from '../lib/store'
+import { StateStore } from '../lib/state_store'
 
 import {
   VpnState,
@@ -12,7 +12,7 @@ import {
   ConnectionState,
 } from '../state/vpn_state'
 
-export function createVpnHandler(store: Store<VpnState>): VpnActions {
+export function createVpnHandler(store: StateStore<VpnState>): VpnActions {
   store.update({
     initialized: true,
     vpnFeatureEnabled: true,