Remember Agentic UI preview path per site when switching sites

apps/ui/src/hooks/use-open-site-url.ts

@@ -45,6 +45,7 @@ export function useOpenSiteUrl( site: SiteDetails ) {
 		try {
 			const redirectTo = new URL( relativeUrl || '/', getSiteUrl( currentSite ) ).toString();
 			preview.setOpen( true );
+			preview.setSite( site.id );
 			preview.updatePath( `/studio-auto-login?redirect_to=${ encodeURIComponent( redirectTo ) }` );
 		} catch ( error ) {
 			// Malformed URL (shouldn't happen) — fall back to the browser so

apps/ui/src/hooks/use-session-ui.test.tsx

@@ -0,0 +1,61 @@
+import { act, renderHook } from '@testing-library/react';
+import { describe, expect, it, vi } from 'vitest';
+import { useConnector } from '@/data/core';
+import { SessionUIProvider, useSessionPreviewUI } from './use-session-ui';
+import type { Connector } from '@/data/core';
+import type { ReactNode } from 'react';
+
+vi.mock( '@/data/core', async ( importOriginal ) => {
+	const actual = await importOriginal< typeof import('@/data/core') >();
+	return {
+		...actual,
+		useConnector: vi.fn(),
+	};
+} );
+
+vi.mocked( useConnector ).mockReturnValue( {
+	onToggleSitePreview: vi.fn( () => vi.fn() ),
+} as unknown as Connector );
+
+function wrapper( { children }: { children: ReactNode } ) {
+	return <SessionUIProvider>{ children }</SessionUIProvider>;
+}
+
+describe( 'useSessionPreviewUI site switching', () => {
+	it( 'opens a never-previewed site at home instead of the previous site path', () => {
+		const { result } = renderHook( () => useSessionPreviewUI(), { wrapper } );
+
+		act( () => result.current.setSite( 'site-a' ) );
+		act( () => result.current.updatePath( '/about' ) );
+		expect( result.current.path ).toBe( '/about' );
+
+		act( () => result.current.setSite( 'site-b' ) );
+		expect( result.current.path ).toBe( '/' );
+		expect( result.current.siteId ).toBe( 'site-b' );
+	} );
+
+	it( 'restores each site last visited path when switching back', () => {
+		const { result } = renderHook( () => useSessionPreviewUI(), { wrapper } );
+
+		act( () => result.current.setSite( 'site-a' ) );
+		act( () => result.current.updatePath( '/about' ) );
+		act( () => result.current.setSite( 'site-b' ) );
+		act( () => result.current.updatePath( '/blog' ) );
+
+		act( () => result.current.setSite( 'site-a' ) );
+		expect( result.current.path ).toBe( '/about' );
+
+		act( () => result.current.setSite( 'site-b' ) );
+		expect( result.current.path ).toBe( '/blog' );
+	} );
+
+	it( 'keeps the path when the previewed site is unchanged', () => {
+		const { result } = renderHook( () => useSessionPreviewUI(), { wrapper } );
+
+		act( () => result.current.setSite( 'site-a' ) );
+		act( () => result.current.updatePath( '/contact' ) );
+
+		act( () => result.current.setSite( 'site-a' ) );
+		expect( result.current.path ).toBe( '/contact' );
+	} );
+} );

apps/ui/src/hooks/use-session-ui.tsx

@@ -27,8 +27,19 @@ import type { Annotation } from '@/components/site-preview/types';
 
 interface PreviewUIState {
 	open: boolean;
-	path: string;
 	reloadNonce: number;
+	// Currently previewed site.
+	siteId: string | null;
+	// Last visited path per site, so switching back restores it. In-memory only.
+	pathsBySiteId: Record< string, string >;
+}
+
+// The path a site's preview shows: its last visited path, or home.
+export function pathForSite(
+	pathsBySiteId: Record< string, string >,
+	siteId: string | null | undefined
+): string {
+	return siteId ? pathsBySiteId[ siteId ] ?? '/' : '/';
 }
 
 export interface SessionUIState {
@@ -40,10 +51,11 @@ export type SessionUIAction =
 	| { type: 'preview/toggle' }
 	| { type: 'preview/navigate'; path: string }
 	| { type: 'preview/reload' }
-	| { type: 'preview/update-path'; path: string };
+	| { type: 'preview/update-path'; path: string }
+	| { type: 'preview/set-site'; siteId: string };
 
 const INITIAL_STATE: SessionUIState = {
-	preview: { open: true, path: '/', reloadNonce: 0 },
+	preview: { open: true, reloadNonce: 0, siteId: null, pathsBySiteId: {} },
 };
 
 function reducer( state: SessionUIState, action: SessionUIAction ): SessionUIState {
@@ -59,7 +71,7 @@ function reducer( state: SessionUIState, action: SessionUIAction ): SessionUISta
 				...state,
 				preview: {
 					...state.preview,
-					path: action.path,
+					pathsBySiteId: rememberPath( state.preview, action.path ),
 					reloadNonce: state.preview.reloadNonce + 1,
 					open: true,
 				},
@@ -75,13 +87,26 @@ function reducer( state: SessionUIState, action: SessionUIAction ): SessionUISta
 					open: true,
 				},
 			};
-		case 'preview/update-path':
-			return state.preview.path === action.path
+		case 'preview/update-path': {
+			const pathsBySiteId = rememberPath( state.preview, action.path );
+			return pathsBySiteId === state.preview.pathsBySiteId
+				? state
+				: { ...state, preview: { ...state.preview, pathsBySiteId } };
+		}
+		case 'preview/set-site':
+			return state.preview.siteId === action.siteId
 				? state
-				: { ...state, preview: { ...state.preview, path: action.path } };
+				: { ...state, preview: { ...state.preview, siteId: action.siteId } };
 	}
 }
 
+function rememberPath( preview: PreviewUIState, path: string ): Record< string, string > {
+	if ( ! preview.siteId || preview.pathsBySiteId[ preview.siteId ] === path ) {
+		return preview.pathsBySiteId;
+	}
+	return { ...preview.pathsBySiteId, [ preview.siteId ]: path };
+}
+
 // Split contexts so that hooks which only need to dispatch (like
 // `useSessionCommands`) don't re-run on every state change.
 const SessionUIStateContext = createContext< SessionUIState | null >( null );
@@ -126,14 +151,6 @@ function SessionUIProviderRoot( { children }: { children: ReactNode } ) {
 	);
 }
 
-function useSessionUIState(): SessionUIState {
-	const value = useContext( SessionUIStateContext );
-	if ( ! value ) {
-		throw new Error( 'useSessionUIState must be used within a SessionUIProvider' );
-	}
-	return value;
-}
-
 export function useSessionUIDispatch(): Dispatch< SessionUIAction > {
 	const value = useContext( SessionUIDispatchContext );
 	if ( ! value ) {
@@ -146,9 +163,12 @@ export interface SessionPreviewUI {
 	readonly open: boolean;
 	readonly path: string;
 	readonly reloadNonce: number;
+	readonly siteId: string | null;
+	readonly pathsBySiteId: Record< string, string >;
 	setOpen: ( value: boolean ) => void;
 	toggle: () => void;
 	updatePath: ( path: string ) => void;
+	setSite: ( siteId: string ) => void;
 }
 
 // Like `useSessionPreviewUI`, but usable outside the dashboard layout —
@@ -166,51 +186,34 @@ export function useOptionalSessionPreviewUI(): SessionPreviewUI | null {
 		( path: string ) => dispatch?.( { type: 'preview/update-path', path } ),
 		[ dispatch ]
 	);
+	const setSite = useCallback(
+		( siteId: string ) => dispatch?.( { type: 'preview/set-site', siteId } ),
+		[ dispatch ]
+	);
 	return useMemo( () => {
 		if ( ! state || ! dispatch ) {
 			return null;
 		}
 		return {
 			open: state.preview.open,
-			path: state.preview.path,
+			path: pathForSite( state.preview.pathsBySiteId, state.preview.siteId ),
 			reloadNonce: state.preview.reloadNonce,
+			siteId: state.preview.siteId,
+			pathsBySiteId: state.preview.pathsBySiteId,
 			setOpen,
 			toggle,
 			updatePath,
+			setSite,
 		};
-	}, [ state, dispatch, setOpen, toggle, updatePath ] );
+	}, [ state, dispatch, setOpen, toggle, updatePath, setSite ] );
 }
 
 export function useSessionPreviewUI(): SessionPreviewUI {
-	const state = useSessionUIState();
-	const dispatch = useSessionUIDispatch();
-	const setOpen = useCallback(
-		( value: boolean ) => dispatch( { type: 'preview/set-open', value } ),
-		[ dispatch ]
-	);
-	const toggle = useCallback( () => dispatch( { type: 'preview/toggle' } ), [ dispatch ] );
-	const updatePath = useCallback(
-		( path: string ) => dispatch( { type: 'preview/update-path', path } ),
-		[ dispatch ]
-	);
-	return useMemo(
-		() => ( {
-			open: state.preview.open,
-			path: state.preview.path,
-			reloadNonce: state.preview.reloadNonce,
-			setOpen,
-			toggle,
-			updatePath,
-		} ),
-		[
-			state.preview.open,
-			state.preview.path,
-			state.preview.reloadNonce,
-			setOpen,
-			toggle,
-			updatePath,
-		]
-	);
+	const ui = useOptionalSessionPreviewUI();
+	if ( ! ui ) {
+		throw new Error( 'useSessionPreviewUI must be used within a SessionUIProvider' );
+	}
+	return ui;
 }
 
 // Registers the on-screen session's annotations handler (its "send to chat")

apps/ui/src/ui-classic/router/layout-dashboard/index.tsx

@@ -10,6 +10,7 @@ import { SitePreview } from '@/components/site-preview';
 import { useSession, useSessionEffectiveEnvironment } from '@/data/queries/use-sessions';
 import { useSites } from '@/data/queries/use-sites';
 import {
+	pathForSite,
 	SessionUIProvider,
 	useSessionPreviewAnnotationsHandler,
 	useSessionPreviewUI,
@@ -104,20 +105,30 @@ function DashboardLayoutContent() {
 		? sites?.find( ( site ) => site.id === lastPreviewSiteId )
 		: undefined;
 	const previewSite = routeSite ?? lastPreviewSite;
+	const previewSiteId = previewSite?.id;
+	const { setSite: setPreviewSite } = preview;
+	useEffect( () => {
+		if ( previewSiteId ) {
+			setPreviewSite( previewSiteId );
+		}
+	}, [ previewSiteId, setPreviewSite ] );
+	// Look up by the route's site so the path is right even before the
+	// `setPreviewSite` effect lands.
+	const previewPath = pathForSite( preview.pathsBySiteId, previewSiteId );
 	const showPreview = preview.open && supportsPreview && !! previewSite;
 	const renderPreview = useCallback(
 		( { collapsed }: PreviewSplitFramePreviewProps ) =>
 			previewSite ? (
 				<SitePreview
 					site={ previewSite }
-					path={ preview.path }
+					path={ previewPath }
 					reloadNonce={ preview.reloadNonce }
 					onAnnotationsDone={ onAnnotationsDone }
 					onPathChange={ preview.updatePath }
 					collapsed={ collapsed }
 				/>
 			) : null,
-		[ onAnnotationsDone, preview.path, preview.reloadNonce, preview.updatePath, previewSite ]
+		[ onAnnotationsDone, previewPath, preview.reloadNonce, preview.updatePath, previewSite ]
 	);
 
 	return (