Studio: Ability to cancel the push and pull

src/hooks/sync-sites/sync-sites-context.tsx

@@ -79,24 +79,22 @@ export function SyncSitesProvider( { children }: { children: React.ReactNode } )
 		[ connectedSites, dispatch ]
 	);
 
-	const { pullSite, isAnySitePulling, isSiteIdPulling, clearPullState, getPullState } = useSyncPull(
-		{
+	const { pullSite, isAnySitePulling, isSiteIdPulling, clearPullState, getPullState, cancelPull } =
+		useSyncPull( {
 			pullStates,
 			setPullStates,
 			onPullSuccess: ( remoteSiteId, localSiteId ) =>
 				updateSiteTimestamp( remoteSiteId, localSiteId, 'pull' ),
-		}
-	);
+		} );
 
 	const [ pushStates, setPushStates ] = useState< PushStates >( {} );
-	const { pushSite, isAnySitePushing, isSiteIdPushing, clearPushState, getPushState } = useSyncPush(
-		{
+	const { pushSite, isAnySitePushing, isSiteIdPushing, clearPushState, getPushState, cancelPush } =
+		useSyncPush( {
 			pushStates,
 			setPushStates,
 			onPushSuccess: ( remoteSiteId, localSiteId ) =>
 				updateSiteTimestamp( remoteSiteId, localSiteId, 'push' ),
-		}
-	);
+		} );
 
 	const { syncSites, isFetching, refetchSites } = useSyncSitesData();
 	useListenDeepLinkConnection( { refetchSites } );
@@ -108,6 +106,7 @@ export function SyncSitesProvider( { children }: { children: React.ReactNode } )
 				isAnySitePulling,
 				isSiteIdPulling,
 				clearPullState,
+				cancelPull,
 				syncSites,
 				refetchSites,
 				isFetching,
@@ -117,6 +116,7 @@ export function SyncSitesProvider( { children }: { children: React.ReactNode } )
 				isAnySitePushing,
 				isSiteIdPushing,
 				clearPushState,
+				cancelPush,
 				getLastSyncTimeText,
 			} }
 		>

src/hooks/sync-sites/use-pull-push-states.ts

@@ -1,4 +1,4 @@
-import { useCallback } from 'react';
+import { useCallback, useRef, useEffect } from 'react';
 
 export const generateStateId = ( selectedSiteId: string, remoteSiteId: number ) =>
 	`${ selectedSiteId }-${ remoteSiteId }`;
@@ -22,31 +22,42 @@ export function usePullPushStates< T >(
 	states: States< T >,
 	setStates: React.Dispatch< React.SetStateAction< States< T > > >
 ): UsePullPushStates< T > {
+	const statesRef = useRef( states );
+
+	useEffect( () => {
+		statesRef.current = states;
+	}, [ states ] );
+
 	const updateState = useCallback< UpdateState< T > >(
 		( selectedSiteId, remoteSiteId, state ) => {
-			setStates( ( prevStates ) => ( {
-				...prevStates,
-				[ generateStateId( selectedSiteId, remoteSiteId ) ]: {
-					...prevStates[ generateStateId( selectedSiteId, remoteSiteId ) ],
-					...state,
-				},
-			} ) );
+			setStates( ( prevStates ) => {
+				const newStates = {
+					...prevStates,
+					[ generateStateId( selectedSiteId, remoteSiteId ) ]: {
+						...prevStates[ generateStateId( selectedSiteId, remoteSiteId ) ],
+						...state,
+					},
+				};
+				statesRef.current = newStates;
+				return newStates;
+			} );
 		},
 		[ setStates ]
 	);
 
 	const getState = useCallback< GetState< T > >(
 		( selectedSiteId, remoteSiteId ): T | undefined => {
-			return states[ generateStateId( selectedSiteId, remoteSiteId ) ];
+			return statesRef.current[ generateStateId( selectedSiteId, remoteSiteId ) ];
 		},
-		[ states ]
+		[]
 	);
 
 	const clearState = useCallback< ClearState >(
 		( selectedSiteId, remoteSiteId ) => {
 			setStates( ( prevStates ) => {
 				const newStates = { ...prevStates };
 				delete newStates[ generateStateId( selectedSiteId, remoteSiteId ) ];
+				statesRef.current = newStates;
 				return newStates;
 			} );
 		},

src/hooks/sync-sites/use-sync-pull.ts

@@ -52,13 +52,16 @@ type UseSyncPullProps = {
 	onPullSuccess?: OnPullSuccess;
 };
 
+type CancelPull = ( selectedSiteId: string, remoteSiteId: number ) => void;
+
 export type UseSyncPull = {
 	pullStates: PullStates;
 	getPullState: GetState< SyncBackupState >;
 	pullSite: PullSite;
 	isAnySitePulling: boolean;
 	isSiteIdPulling: IsSiteIdPulling;
 	clearPullState: ClearState;
+	cancelPull: CancelPull;
 };
 
 export function useSyncPull( {
@@ -74,6 +77,7 @@ export function useSyncPull( {
 		isKeyPulling,
 		isKeyFinished,
 		isKeyFailed,
+		isKeyCancelled,
 		getBackupStatusWithProgress,
 	} = useSyncStatesProgressInfo();
 	const {
@@ -87,13 +91,13 @@ export function useSyncPull( {
 			updateState( selectedSiteId, remoteSiteId, state );
 			const statusKey = state.status?.key;
 
-			if ( isKeyFailed( statusKey ) || isKeyFinished( statusKey ) ) {
+			if ( isKeyFailed( statusKey ) || isKeyFinished( statusKey ) || isKeyCancelled( statusKey ) ) {
 				getIpcApi().clearSyncOperation( generateStateId( selectedSiteId, remoteSiteId ) );
 			} else {
 				getIpcApi().addSyncOperation( generateStateId( selectedSiteId, remoteSiteId ) );
 			}
 		},
-		[ isKeyFailed, isKeyFinished, updateState ]
+		[ isKeyFailed, isKeyFinished, isKeyCancelled, updateState ]
 	);
 
 	const clearPullState = useCallback< ClearState >(
@@ -114,6 +118,8 @@ export function useSyncPull( {
 
 			const remoteSiteId = connectedSite.id;
 			const remoteSiteUrl = connectedSite.url;
+
+			clearPullState( selectedSite.id, remoteSiteId );
 			updatePullState( selectedSite.id, remoteSiteId, {
 				backupId: null,
 				status: pullStatesProgressInfo[ 'in-progress' ],
@@ -161,7 +167,7 @@ export function useSyncPull( {
 				} );
 			}
 		},
-		[ __, client, pullStatesProgressInfo, updatePullState ]
+		[ __, clearPullState, client, pullStatesProgressInfo, updatePullState ]
 	);
 
 	const checkBackupFileSize = async ( downloadUrl: string ): Promise< number > => {
@@ -213,7 +219,17 @@ export function useSyncPull( {
 					downloadUrl,
 				} );
 
-				const filePath = await getIpcApi().downloadSyncBackup( remoteSiteId, downloadUrl );
+				const operationId = generateStateId( selectedSite.id, remoteSiteId );
+				const filePath = await getIpcApi().downloadSyncBackup(
+					remoteSiteId,
+					downloadUrl,
+					operationId
+				);
+
+				const stateAfterDownload = getPullState( selectedSite.id, remoteSiteId );
+				if ( ! stateAfterDownload || isKeyCancelled( stateAfterDownload?.status.key ) ) {
+					return;
+				}
 
 				// Starting import process
 				updatePullState( selectedSite.id, remoteSiteId, {
@@ -252,6 +268,12 @@ export function useSyncPull( {
 				onPullSuccess?.( remoteSiteId, selectedSite.id );
 			} catch ( error ) {
 				console.error( 'Backup completion failed:', error );
+
+				const currentState = getPullState( selectedSite.id, remoteSiteId );
+				if ( currentState && isKeyCancelled( currentState?.status.key ) ) {
+					return;
+				}
+
 				Sentry.captureException( error );
 				updatePullState( selectedSite.id, remoteSiteId, {
 					status: pullStatesProgressInfo.failed,
@@ -266,8 +288,10 @@ export function useSyncPull( {
 			__,
 			clearImportState,
 			clearPullState,
+			getPullState,
 			importFile,
 			onPullSuccess,
+			isKeyCancelled,
 			pullStatesProgressInfo.cancelled,
 			pullStatesProgressInfo.downloading,
 			pullStatesProgressInfo.failed,
@@ -284,7 +308,12 @@ export function useSyncPull( {
 				return;
 			}
 
-			const backupId = getPullState( selectedSiteId, remoteSiteId )?.backupId;
+			const currentState = getPullState( selectedSiteId, remoteSiteId );
+			if ( currentState && isKeyCancelled( currentState.status.key ) ) {
+				return;
+			}
+
+			const backupId = currentState?.backupId;
 			if ( ! backupId ) {
 				console.error( 'No backup ID found' );
 				return;
@@ -335,13 +364,18 @@ export function useSyncPull( {
 			onBackupCompleted,
 			pullStatesProgressInfo,
 			updatePullState,
+			isKeyCancelled,
 		]
 	);
 
 	useEffect( () => {
 		const intervals: Record< string, NodeJS.Timeout > = {};
 
 		Object.entries( pullStates ).forEach( ( [ key, state ] ) => {
+			if ( isKeyCancelled( state.status.key ) ) {
+				return;
+			}
+
 			if ( state.backupId && state.status.key === 'in-progress' ) {
 				intervals[ key ] = setTimeout( () => {
 					void fetchAndUpdateBackup( state.remoteSiteId, state.selectedSite.id );
@@ -352,7 +386,7 @@ export function useSyncPull( {
 		return () => {
 			Object.values( intervals ).forEach( clearTimeout );
 		};
-	}, [ pullStates, fetchAndUpdateBackup ] );
+	}, [ pullStates, fetchAndUpdateBackup, isKeyCancelled ] );
 
 	const isAnySitePulling = useMemo< boolean >( () => {
 		return Object.values( pullStates ).some( ( state ) => isKeyPulling( state.status.key ) );
@@ -361,6 +395,9 @@ export function useSyncPull( {
 	const isSiteIdPulling = useCallback< IsSiteIdPulling >(
 		( selectedSiteId, remoteSiteId ) => {
 			return Object.values( pullStates ).some( ( state ) => {
+				if ( ! state.selectedSite ) {
+					return false;
+				}
 				if ( state.selectedSite.id !== selectedSiteId ) {
 					return false;
 				}
@@ -373,5 +410,37 @@ export function useSyncPull( {
 		[ pullStates, isKeyPulling ]
 	);
 
-	return { pullStates, getPullState, pullSite, isAnySitePulling, isSiteIdPulling, clearPullState };
+	const cancelPull = useCallback< CancelPull >(
+		async ( selectedSiteId, remoteSiteId ) => {
+			const operationId = generateStateId( selectedSiteId, remoteSiteId );
+
+			await getIpcApi().cancelSyncOperation( operationId );
+
+			updatePullState( selectedSiteId, remoteSiteId, {
+				status: pullStatesProgressInfo.cancelled,
+			} );
+
+			getIpcApi()
+				.removeSyncBackup( remoteSiteId )
+				.catch( () => {
+					// Ignore errors if file doesn't exist
+				} );
+
+			getIpcApi().showNotification( {
+				title: __( 'Pull cancelled' ),
+				body: __( 'The pull operation has been cancelled.' ),
+			} );
+		},
+		[ __, pullStatesProgressInfo.cancelled, updatePullState ]
+	);
+
+	return {
+		pullStates,
+		getPullState,
+		pullSite,
+		isAnySitePulling,
+		isSiteIdPulling,
+		clearPullState,
+		cancelPull,
+	};
 }