Fix theme previewing

packages/edit-site/src/components/app/index.js

@@ -11,6 +11,7 @@ import { useDispatch, useSelect } from '@wordpress/data';
 import { __, sprintf } from '@wordpress/i18n';
 import { PluginArea } from '@wordpress/plugins';
 import { privateApis as routerPrivateApis } from '@wordpress/router';
+import { useMemo } from '@wordpress/element';
 
 /**
  * Internal dependencies
@@ -21,6 +22,10 @@ import { store as editSiteStore } from '../../store';
 import { useCommonCommands } from '../../hooks/commands/use-common-commands';
 import useSetCommandContext from '../../hooks/commands/use-set-command-context';
 import { useRegisterSiteEditorRoutes } from '../site-editor-routes';
+import {
+	currentlyPreviewingTheme,
+	isPreviewingTheme,
+} from '../../utils/is-previewing-theme';
 
 const { RouterProvider } = unlock( routerPrivateApis );
 const { GlobalStylesProvider } = unlock( editorPrivateApis );
@@ -49,12 +54,37 @@ export default function App() {
 			)
 		);
 	}
+	const middlewares = useMemo(
+		() => [
+			( { path, query } ) => {
+				if ( ! isPreviewingTheme() ) {
+					return { path, query };
+				}
+
+				return {
+					path,
+					query: {
+						...query,
+						wp_theme_preview:
+							'wp_theme_preview' in query
+								? query.wp_theme_preview
+								: currentlyPreviewingTheme(),
+					},
+				};
+			},
+		],
+		[]
+	);
 
 	return (
 		<SlotFillProvider>
 			<GlobalStylesProvider>
 				<UnsavedChangesWarning />
-				<RouterProvider routes={ routes } pathArg="p">
+				<RouterProvider
+					routes={ routes }
+					pathArg="p"
+					middlewares={ middlewares }
+				>
 					<AppLayout />
 					<PluginArea onError={ onPluginAreaError } />
 				</RouterProvider>

packages/edit-site/src/utils/is-previewing-theme.js

@@ -4,9 +4,7 @@
 import { getQueryArg } from '@wordpress/url';
 
 export function isPreviewingTheme() {
-	return (
-		getQueryArg( window.location.href, 'wp_theme_preview' ) !== undefined
-	);
+	return !! getQueryArg( window.location.href, 'wp_theme_preview' );
 }
 
 export function currentlyPreviewingTheme() {

packages/edit-site/src/utils/use-activate-theme.js

@@ -4,6 +4,7 @@
 import { store as coreStore } from '@wordpress/core-data';
 import { useDispatch } from '@wordpress/data';
 import { privateApis as routerPrivateApis } from '@wordpress/router';
+import { addQueryArgs } from '@wordpress/url';
 
 /**
  * Internal dependencies
@@ -14,7 +15,7 @@ import {
 	currentlyPreviewingTheme,
 } from './is-previewing-theme';
 
-const { useHistory } = unlock( routerPrivateApis );
+const { useHistory, useLocation } = unlock( routerPrivateApis );
 
 /**
  * This should be refactored to use the REST API, once the REST API can activate themes.
@@ -23,6 +24,7 @@ const { useHistory } = unlock( routerPrivateApis );
  */
 export function useActivateTheme() {
 	const history = useHistory();
+	const { path } = useLocation();
 	const { startResolution, finishResolution } = useDispatch( coreStore );
 
 	return async () => {
@@ -37,8 +39,7 @@ export function useActivateTheme() {
 			finishResolution( 'activateTheme' );
 			// Remove the wp_theme_preview query param: we've finished activating
 			// the queue and are switching to normal Site Editor.
-			const { params } = history.getLocationWithParams();
-			history.replace( { ...params, wp_theme_preview: undefined } );
+			history.navigate( addQueryArgs( path, { wp_theme_preview: '' } ) );
 		}
 	};
 }

packages/router/package.json

@@ -29,6 +29,7 @@
 	"types": "build-types",
 	"dependencies": {
 		"@babel/runtime": "7.25.7",
+		"@wordpress/compose": "*",
 		"@wordpress/element": "*",
 		"@wordpress/private-apis": "*",
 		"@wordpress/url": "*",

packages/router/src/link.tsx

@@ -1,30 +1,42 @@
 /**
  * WordPress dependencies
  */
-import { useContext } from '@wordpress/element';
+import { useContext, useMemo } from '@wordpress/element';
 import { getQueryArgs, getPath, buildQueryString } from '@wordpress/url';
+import { compose } from '@wordpress/compose';
 
 /**
  * Internal dependencies
  */
-import { ConfigContext, type NavigationOptions, useHistory } from './router';
+import {
+	ConfigContext,
+	type Middleware,
+	type NavigationOptions,
+	useHistory,
+} from './router';
 
 export function useLink( to: string, options: NavigationOptions = {} ) {
 	const history = useHistory();
-	const { pathArg } = useContext( ConfigContext );
+	const { pathArg, middlewares } = useContext( ConfigContext );
 	function onClick( event: React.SyntheticEvent< HTMLAnchorElement > ) {
 		event?.preventDefault();
 		history.navigate( to, options );
 	}
-	const queryArgs = getQueryArgs( to );
-	const path = getPath( 'http://domain.com/' + to );
+	const query = getQueryArgs( to );
+	const path = getPath( 'http://domain.com/' + to ) ?? '';
+	const link = useMemo( () => {
+		const runMiddlewares = (
+			middlewares ? compose( ...middlewares ) : ( i: unknown ) => i
+		) as Middleware;
+		return runMiddlewares( { path, query } );
+	}, [ path, query, middlewares ] );
 
 	const [ before ] = window.location.href.split( '?' );
 
 	return {
 		href: `${ before }?${ buildQueryString( {
-			[ pathArg ]: path,
-			...queryArgs,
+			[ pathArg ]: link.path,
+			...link.query,
 		} ) }`,
 		onClick,
 	};

packages/router/src/router.tsx

@@ -19,6 +19,7 @@ import {
 	getPath,
 	buildQueryString,
 } from '@wordpress/url';
+import { compose } from '@wordpress/compose';
 
 /**
  * Internal dependencies
@@ -46,8 +47,17 @@ interface Match {
 	params?: Record< string, any >;
 }
 
+export type Middleware = ( arg: {
+	path: string;
+	query: Record< string, any >;
+} ) => {
+	path: string;
+	query: Record< string, any >;
+};
+
 interface Config {
 	pathArg: string;
+	middlewares?: Middleware[];
 }
 
 export interface NavigationOptions {
@@ -77,18 +87,24 @@ export function useLocation() {
 }
 
 export function useHistory() {
-	const { pathArg } = useContext( ConfigContext );
+	const { pathArg, middlewares } = useContext( ConfigContext );
 	return useMemo(
 		() => ( {
 			navigate( rawPath: string, options: NavigationOptions = {} ) {
+				const runMiddlewares = (
+					middlewares
+						? compose( ...middlewares )
+						: ( i: unknown ) => i
+				) as Middleware;
 				const query = getQueryArgs( rawPath );
-				const path = getPath( 'http://domain.com/' + rawPath );
+				const path = getPath( 'http://domain.com/' + rawPath ) ?? '';
 				const performPush = () => {
+					const result = runMiddlewares( { path, query } );
 					return history.push(
 						{
 							search: buildQueryString( {
-								[ pathArg ]: path,
-								...query,
+								[ pathArg ]: result.path,
+								...result.query,
 							} ),
 						},
 						options.state
@@ -121,7 +137,7 @@ export function useHistory() {
 				} );
 			},
 		} ),
-		[ pathArg ]
+		[ pathArg, middlewares ]
 	);
 }
 
@@ -180,10 +196,12 @@ export default function useMatch(
 export function RouterProvider( {
 	routes,
 	pathArg,
+	middlewares,
 	children,
 }: {
 	routes: Route[];
 	pathArg: string;
+	middlewares?: Middleware[];
 	children: React.ReactNode;
 } ) {
 	const location = useSyncExternalStore(
@@ -192,7 +210,10 @@ export function RouterProvider( {
 		getLocationWithQuery
 	);
 	const match = useMatch( location, routes, pathArg );
-	const config = useMemo( () => ( { pathArg } ), [ pathArg ] );
+	const config = useMemo(
+		() => ( { middlewares, pathArg } ),
+		[ middlewares, pathArg ]
+	);
 
 	return (
 		<ConfigContext.Provider value={ config }>

packages/router/tsconfig.json

@@ -4,11 +4,10 @@
 	"compilerOptions": {
 		"rootDir": "src",
 		"declarationDir": "build-types",
-		"types": [ "gutenberg-env" ],
-		"allowJs": false,
-		"checkJs": false
+		"types": [ "gutenberg-env" ]
 	},
 	"references": [
+		{ "path": "../compose" },
 		{ "path": "../element" },
 		{ "path": "../private-apis" },
 		{ "path": "../url" }