React Hooks to save/restore arbitrary object bound to the history stack of Next.js, to restore view state incl. scroll position on go back or forward.

useHistorySavedState.tsx

import { HistoryState } from 'next/dist/shared/lib/router/router'
import { useRouter } from 'next/router'
import React, { useContext, useRef } from 'react'
import { useIsomorphicLayoutEffect } from '...' // please place it here yourself or import it from your project.

interface ManagerState {
  // matches to history.state
  restoredKey?: string
  bagForRestore: Record<string, unknown>

  // update is delayed until routeChangeComplete, as popstate is triggered after history.state is changed
  renderedState?: HistoryState

  saveCallbackMap: Map<string, () => unknown>
}

const identifier = "yourProductName"

function useHistorySavedStateManager() {
  const router = useRouter()
  const stateRef = useRef<ManagerState>({
    bagForRestore: {},
    saveCallbackMap: new Map(),
  })

  useIsomorphicLayoutEffect(() => {
    const handleEnd = () => {
      stateRef.current.renderedState = window.history.state
    }

    const handleSaveState = () => {
      const state = stateRef.current.renderedState
      const key = state?.__N && state.key
      if (!key) return

      const bag: Record<string, unknown> = {}
      stateRef.current.saveCallbackMap.forEach((save, key) => {
        try {
          bag[key] = save()
        } catch (e: unknown) {
          console.error('suppressed error while save state: ', e)
        }
      })
      sessionStorage.setItem(`${identifier}.historySavedState.${key}`, JSON.stringify(bag))
      if (key === stateRef.current.restoredKey) {
        stateRef.current.bagForRestore = bag
      }
    }

    const handleVisibilityChange = () => {
      if (document.visibilityState === 'hidden') {
        handleSaveState()
      }
    }

    router.events.on('routeChangeComplete', handleEnd)
    router.events.on('beforeHistoryChange', handleSaveState)
    // called when tab is closed or just user switch to another tab
    // window.addEventListener('beforeunload', handleSaveState)
    window.addEventListener('visibilitychange', handleVisibilityChange)
    return () => {
      router.events.off('routeChangeComplete', handleEnd)
      router.events.off('beforeHistoryChange', handleSaveState)
      // window.removeEventListener('beforeunload', handleSaveState)
      window.removeEventListener('visibilitychange', handleVisibilityChange)
    }
  }, [router.events])

  return stateRef
}

const StateRefContext = React.createContext<React.MutableRefObject<ManagerState> | undefined>(undefined)

export function NextHistoryStateProvider({ children }: { children: React.ReactNode }) {
  const stateRef = useHistorySavedStateManager()

  return <StateRefContext.Provider value={stateRef}>{children}</StateRefContext.Provider>
}

// We does not "return" saved state (client-only state)  because it causes hydration mismatch error on the initial render.
export function useHistorySavedState<T>(
  key: string | null,
  handlers: {
    onRestoreState: (savedState: T) => void
    onSaveState: () => T
  }
): void {
  if (key === '') throw new Error('key must be non-empty string or null')
  const stateRef = useContext(StateRefContext)
  if (!stateRef) {
    throw new Error('HistorySaveStateManager not provided in tree.')
  }

  const handlersRef = useRef(handlers)
  useIsomorphicLayoutEffect(() => {
    handlersRef.current = handlers
  }, [handlers])

  useIsomorphicLayoutEffect(() => {
    if (!key) return

    // There is no hook for "after change state" (popstate is for browser back/forward only) so lazily read restore bag here.
    const state = window.history.state
    if (!state) return
    const stateKey = state.key as string | undefined
    if (!stateKey) {
      if (process.env.NODE_ENV === 'development') {
        throw new Error('history.state.key is not set by Next.js')
      }
      return
    }

    if (stateKey !== stateRef.current.restoredKey) {
      try {
        stateRef.current.bagForRestore =
          JSON.parse(sessionStorage.getItem(`${identifier}.historySavedState.${stateKey}`) ?? 'null') ?? {}
      } catch (e: unknown) {
        console.warn('Failed to restore from history state', e)
        stateRef.current.bagForRestore = {}
      }
      stateRef.current.restoredKey = stateKey
    }

    const restoredState = stateRef.current.bagForRestore[key]
    if (restoredState) {
      handlersRef.current.onRestoreState(restoredState as T)
    }

    if (stateRef.current.saveCallbackMap.has(key)) {
      if (process.env.NODE_ENV === 'development') {
        throw new Error(`history saved state key "${key}" is already registered`)
      }
      return
    }

    stateRef.current.saveCallbackMap.set(key, () => handlersRef.current.onSaveState())

    return () => {
      stateRef.current.saveCallbackMap.delete(key)
    }
  }, [key])
}