MSD: read origin_site_id param and set as current omnibar site

client/dashboard/app/interim-omnibar/index.tsx

@@ -5,8 +5,10 @@ import {
 	userPreferenceQuery,
 } from '@automattic/api-queries';
 import { QueryObserver } from '@tanstack/react-query';
+import { removeQueryArgs } from '@wordpress/url';
 import { hydrateRoot } from 'react-dom/client';
 import { AUTH_QUERY_KEY, initializeCurrentUser } from '../auth';
+import { getCurrentOmnibarSiteId, setCurrentOmnibarSiteId } from '../omnibar/current-site';
 import type { OmnibarEvents } from './click-handlers';
 import type { UserPreferences } from '@automattic/api-core';
 
@@ -56,7 +58,9 @@ export default async function loadOmnibar( events: OmnibarEvents ) {
 		);
 
 	async function renderWithSiteId( siteId: number | undefined ) {
-		const site = siteId ? await queryClient.ensureQueryData( siteByIdQuery( siteId ) ) : null;
+		const site = siteId
+			? await queryClient.ensureQueryData( siteByIdQuery( siteId ) ).catch( () => null )
+			: null;
 		root.render(
 			<InterimOmnibar
 				user={ user }
@@ -68,19 +72,30 @@ export default async function loadOmnibar( events: OmnibarEvents ) {
 		);
 	}
 
-	// Render with the initial recent site (or primary blog as fallback).
+	const originSiteId = Number(
+		new URLSearchParams( window.location.search ).get( 'origin_site_id' )
+	);
+	if ( originSiteId > 0 ) {
+		setCurrentOmnibarSiteId( originSiteId );
+		window.history.replaceState(
+			null,
+			'',
+			removeQueryArgs( window.location.pathname + window.location.search, 'origin_site_id' )
+		);
+	}
+
 	const recentSites = queryClient.getQueryData< UserPreferences >(
 		rawUserPreferencesQuery().queryKey
 	)?.recentSites;
-	const initialSiteId = recentSites?.[ 0 ] || user.primary_blog;
+
+	// Render with the origin site as higher priority.
+	const initialSiteId = originSiteId || getCurrentOmnibarSiteId( user, recentSites );
 	renderWithSiteId( initialSiteId );
 
 	// Re-render whenever recentSites changes (e.g. user navigates to a different site).
-	let currentSiteId = initialSiteId;
 	new QueryObserver( queryClient, userPreferenceQuery( 'recentSites' ) ).subscribe( ( result ) => {
-		const siteId = result.data?.[ 0 ] || user.primary_blog;
-		if ( siteId !== currentSiteId ) {
-			currentSiteId = siteId;
+		const siteId = getCurrentOmnibarSiteId( user, result.data );
+		if ( siteId !== initialSiteId ) {
 			renderWithSiteId( siteId );
 		}
 	} );

client/dashboard/app/omnibar/current-site.ts

@@ -0,0 +1,31 @@
+import {
+	queryClient,
+	rawUserPreferencesQuery,
+	userPreferenceOptimisticMutation,
+} from '@automattic/api-queries';
+import { MutationObserver } from '@tanstack/react-query';
+import type { User } from '@automattic/api-core';
+
+/**
+ * Returns the current site ID to be displayed in the omnibar,
+ * based on the user's recent sites and primary blog.
+ */
+export function getCurrentOmnibarSiteId( user: User, recentSites?: number[] ) {
+	return recentSites?.[ 0 ] || user.primary_blog;
+}
+
+/**
+ * Sets the current site to be displayed in the omnibar,
+ * by pushing the site to the front of the user's recent sites preferences.
+ */
+export async function setCurrentOmnibarSiteId( siteId: number ) {
+	const prefs = await queryClient.ensureQueryData( rawUserPreferencesQuery() );
+	const recentSites = prefs?.recentSites ?? [];
+	if ( siteId === recentSites[ 0 ] ) {
+		return;
+	}
+	const updated = [ ...new Set( [ siteId, ...recentSites ] ) ].slice( 0, 5 );
+	new MutationObserver( queryClient, userPreferenceOptimisticMutation( 'recentSites' ) ).mutate(
+		updated
+	);
+}

client/dashboard/app/router/sites.tsx

@@ -36,13 +36,11 @@ import {
 	siteSshAccessStatusQuery,
 	siteStaticFile404SettingQuery,
 	siteWordPressVersionQuery,
-	userPreferenceOptimisticMutation,
 	queryClient,
 	wpOrgCoreVersionQuery,
 } from '@automattic/api-queries';
 import { isEnabled } from '@automattic/calypso-config';
 import { isSupportSession } from '@automattic/calypso-support-session';
-import { MutationObserver } from '@tanstack/react-query';
 import {
 	createLazyRoute,
 	createRoute,
@@ -69,6 +67,7 @@ import { hasSiteTrialEnded } from '../../utils/site-trial';
 import { getSiteTypeFeatureSupports } from '../../utils/site-type-feature-support';
 import { isSelfHostedJetpackConnected } from '../../utils/site-types';
 import { AUTH_QUERY_KEY } from '../auth';
+import { setCurrentOmnibarSiteId } from '../omnibar/current-site';
 import { rootRoute } from './root';
 import type { AppConfig } from '../context';
 import type { DifmWebsiteContentResponse, Site, User } from '@automattic/api-core';
@@ -173,16 +172,8 @@ export const siteRoute = createRoute( {
 	},
 	onEnter: async ( { loaderData } ) => {
 		const siteId = loaderData?.site?.ID;
-		if ( ! siteId ) {
-			return;
-		}
-		const prefs = await queryClient.ensureQueryData( rawUserPreferencesQuery() );
-		const recentSites = prefs?.recentSites ?? [];
-		if ( siteId !== recentSites[ 0 ] ) {
-			const updated = [ ...new Set( [ siteId, ...recentSites ] ) ].slice( 0, 5 );
-			new MutationObserver( queryClient, userPreferenceOptimisticMutation( 'recentSites' ) ).mutate(
-				updated
-			);
+		if ( siteId ) {
+			setCurrentOmnibarSiteId( siteId );
 		}
 	},
 	errorComponent: lazyRouteComponent( () => import( '../../sites/site/error' ) ),