Clean up WordPress.com connections for deleted sites

apps/cli/commands/site/delete.ts

@@ -1,6 +1,7 @@
 import fs from 'fs';
 import { deleteAiSessionsForSite } from '@studio/common/ai/sessions/manage';
 import { SITE_EVENTS } from '@studio/common/lib/cli-events';
+import { removeAllConnectedWpcomSitesForLocalSite } from '@studio/common/lib/connected-sites';
 import { arePathsEqual } from '@studio/common/lib/fs-utils';
 import { readAuthToken, type StoredAuthToken } from '@studio/common/lib/shared-config';
 import { getSessionsDirectory } from '@studio/common/lib/well-known-paths';
@@ -118,6 +119,18 @@ export async function runCommand(
 			await unlockCliConfig();
 		}
 
+		try {
+			await removeAllConnectedWpcomSitesForLocalSite( site.id );
+		} catch ( error ) {
+			logger.reportError(
+				new LoggerError(
+					__( 'Failed to remove WordPress.com connections. Proceeding anyway…' ),
+					error
+				),
+				false
+			);
+		}
+
 		try {
 			await deleteAiSessionsForSite( getSessionsDirectory(), {
 				id: site.id,

apps/cli/commands/site/tests/delete.test.ts

@@ -1,5 +1,6 @@
 import fs from 'fs';
 import { deleteAiSessionsForSite } from '@studio/common/ai/sessions/manage';
+import { removeAllConnectedWpcomSitesForLocalSite } from '@studio/common/lib/connected-sites';
 import { arePathsEqual } from '@studio/common/lib/fs-utils';
 import { readAuthToken } from '@studio/common/lib/shared-config';
 import { SITE_RUNTIME_PLAYGROUND } from '@studio/common/lib/site-runtime';
@@ -26,6 +27,9 @@ import { runCommand } from '../delete';
 vi.mock( 'fs' );
 vi.mock( 'cli/lib/api' );
 vi.mock( '@studio/common/ai/sessions/manage' );
+vi.mock( '@studio/common/lib/connected-sites', () => ( {
+	removeAllConnectedWpcomSitesForLocalSite: vi.fn(),
+} ) );
 vi.mock( import( '@studio/common/lib/shared-config' ), async ( importOriginal ) => ( {
 	...( await importOriginal() ),
 	readAuthToken: vi.fn(),
@@ -134,6 +138,7 @@ describe( 'CLI: studio site delete', () => {
 		vi.mocked( stopProxyIfNoSitesNeedIt ).mockResolvedValue( undefined );
 		vi.mocked( arePathsEqual ).mockImplementation( ( a: string, b: string ) => a === b );
 		vi.mocked( deleteAiSessionsForSite ).mockResolvedValue( [] );
+		vi.mocked( removeAllConnectedWpcomSitesForLocalSite ).mockResolvedValue( undefined );
 		vi.spyOn( fs, 'existsSync' ).mockReturnValue( true );
 	} );
 
@@ -208,6 +213,7 @@ describe( 'CLI: studio site delete', () => {
 			expect( lockCliConfig ).toHaveBeenCalled();
 			expect( readCliConfig ).toHaveBeenCalled();
 			expect( saveCliConfig ).toHaveBeenCalled();
+			expect( removeAllConnectedWpcomSitesForLocalSite ).toHaveBeenCalledWith( testSite.id );
 			const savedCliConfig = vi.mocked( saveCliConfig ).mock.calls[ 0 ][ 0 ];
 			expect( savedCliConfig.sites ).toHaveLength( 0 );
 			expect( unlockCliConfig ).toHaveBeenCalled();
@@ -365,6 +371,17 @@ describe( 'CLI: studio site delete', () => {
 			expect( disconnectFromDaemon ).toHaveBeenCalled();
 		} );
 
+		it( 'should proceed when deleting WordPress.com connections fails', async () => {
+			vi.mocked( removeAllConnectedWpcomSitesForLocalSite ).mockRejectedValue(
+				new Error( 'shared config failed' )
+			);
+
+			await expect( runCommand( testSiteFolder, true ) ).resolves.not.toThrow();
+			expect( saveCliConfig ).toHaveBeenCalled();
+			expect( trash ).toHaveBeenCalledWith( [ testSiteFolder ] );
+			expect( disconnectFromDaemon ).toHaveBeenCalled();
+		} );
+
 		it( 'should not remove domain or certificate if no custom domain', async () => {
 			vi.mocked( getSnapshotsFromConfig ).mockResolvedValue( [] );
 

apps/cli/migrations/07-cleanup-orphaned-connected-sites.ts

@@ -0,0 +1,54 @@
+import {
+	lockSharedConfig,
+	readSharedConfig,
+	saveSharedConfig,
+	unlockSharedConfig,
+} from '@studio/common/lib/shared-config';
+import { readCliConfig } from 'cli/lib/cli-config/core';
+import type { Migration } from '@studio/common/lib/migration';
+import type { SharedConfig } from '@studio/common/lib/shared-config';
+
+function hasOrphanedConnections( config: SharedConfig, localSiteIds: Set< string > ): boolean {
+	return Object.values( config.connectedWpcomSites ?? {} ).some( ( connections ) =>
+		connections.some( ( connection ) => ! localSiteIds.has( connection.localSiteId ) )
+	);
+}
+
+async function getLocalSiteIds(): Promise< Set< string > > {
+	return new Set( ( await readCliConfig() ).sites.map( ( site ) => site.id ) );
+}
+
+export const cleanupOrphanedConnectedSites: Migration = {
+	async needsToRun() {
+		const [ config, localSiteIds ] = await Promise.all( [ readSharedConfig(), getLocalSiteIds() ] );
+		return hasOrphanedConnections( config, localSiteIds );
+	},
+
+	async run() {
+		try {
+			await lockSharedConfig();
+			const config = await readSharedConfig();
+			const localSiteIds = await getLocalSiteIds();
+			if ( ! hasOrphanedConnections( config, localSiteIds ) ) {
+				return;
+			}
+
+			for ( const [ userId, connections ] of Object.entries( config.connectedWpcomSites ?? {} ) ) {
+				const validConnections = connections.filter( ( connection ) =>
+					localSiteIds.has( connection.localSiteId )
+				);
+				if ( validConnections.length > 0 ) {
+					config.connectedWpcomSites![ userId ] = validConnections;
+				} else {
+					delete config.connectedWpcomSites![ userId ];
+				}
+			}
+			if ( config.connectedWpcomSites && Object.keys( config.connectedWpcomSites ).length === 0 ) {
+				delete config.connectedWpcomSites;
+			}
+			await saveSharedConfig( config );
+		} finally {
+			await unlockSharedConfig();
+		}
+	},
+};

apps/cli/migrations/index.ts

@@ -5,6 +5,7 @@ import { renameProcessManagerHome } from './03-rename-pm-home';
 import { cleanupObsoleteServerFiles } from './04-cleanup-obsolete-server-files';
 import { migrateConnectedSitesToShared } from './05-migrate-connected-sites-to-shared';
 import { installBundledDefaultPhp } from './06-install-bundled-default-php';
+import { cleanupOrphanedConnectedSites } from './07-cleanup-orphaned-connected-sites';
 import type { Migration } from '@studio/common/lib/migration';
 
 export const migrations: Migration[] = [
@@ -14,5 +15,6 @@ export const migrations: Migration[] = [
 	cleanupObsoleteServerFiles,
 	migrateConnectedSitesToShared,
 	installBundledDefaultPhp,
+	cleanupOrphanedConnectedSites,
 	moveAiSessionsToStudioDir,
 ];

apps/cli/migrations/tests/07-cleanup-orphaned-connected-sites.test.ts

@@ -0,0 +1,75 @@
+import {
+	lockSharedConfig,
+	readSharedConfig,
+	saveSharedConfig,
+	unlockSharedConfig,
+} from '@studio/common/lib/shared-config';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { readCliConfig } from 'cli/lib/cli-config/core';
+import { cleanupOrphanedConnectedSites } from '../07-cleanup-orphaned-connected-sites';
+import type { SyncSite } from '@studio/common/types/sync';
+
+vi.mock( '@studio/common/lib/shared-config', () => ( {
+	lockSharedConfig: vi.fn(),
+	readSharedConfig: vi.fn(),
+	saveSharedConfig: vi.fn(),
+	unlockSharedConfig: vi.fn(),
+} ) );
+vi.mock( 'cli/lib/cli-config/core', () => ( {
+	readCliConfig: vi.fn(),
+} ) );
+
+function connection( id: number, localSiteId: string ): SyncSite {
+	return {
+		id,
+		localSiteId,
+		name: `Site ${ id }`,
+		url: `https://site-${ id }.example.com`,
+		isStaging: false,
+		isPressable: false,
+		syncSupport: 'already-connected',
+		lastPullTimestamp: null,
+		lastPushTimestamp: null,
+	};
+}
+
+describe( 'cleanupOrphanedConnectedSites', () => {
+	beforeEach( () => {
+		vi.clearAllMocks();
+		vi.mocked( readCliConfig, { partial: true } ).mockResolvedValue( {
+			version: 1,
+			sites: [
+				{
+					id: 'local-a',
+					name: 'Local Site',
+					path: '/local-a',
+					port: 8881,
+					phpVersion: '8.0',
+				},
+			],
+			snapshots: [],
+		} );
+	} );
+
+	it( 'removes orphaned connections across account buckets while preserving valid ones', async () => {
+		const config = {
+			version: 1 as const,
+			connectedWpcomSites: {
+				'7': [ connection( 1, 'local-a' ), connection( 2, 'deleted-site' ) ],
+				'8': [ connection( 3, 'deleted-site' ) ],
+			},
+		};
+		vi.mocked( readSharedConfig ).mockResolvedValue( config );
+
+		await expect( cleanupOrphanedConnectedSites.needsToRun() ).resolves.toBe( true );
+		await cleanupOrphanedConnectedSites.run();
+
+		expect( lockSharedConfig ).toHaveBeenCalledOnce();
+		expect( saveSharedConfig ).toHaveBeenCalledWith( config );
+		expect( unlockSharedConfig ).toHaveBeenCalledOnce();
+		expect( config.connectedWpcomSites ).toEqual( {
+			'7': [ connection( 1, 'local-a' ) ],
+		} );
+		await expect( cleanupOrphanedConnectedSites.needsToRun() ).resolves.toBe( false );
+	} );
+} );

packages/common/lib/connected-sites.ts

@@ -198,6 +198,25 @@ export async function removeConnectedWpcomSite(
 	}
 }
 
+export async function removeAllConnectedWpcomSitesForLocalSite(
+	localSiteId: string
+): Promise< void > {
+	try {
+		await lockSharedConfig();
+		const config = await readSharedConfig();
+		for ( const userKey of Object.keys( config.connectedWpcomSites ?? {} ) ) {
+			const connections = getConnectionsForUser( config, userKey );
+			config.connectedWpcomSites![ userKey ] = connections.filter(
+				( connection ) => connection.localSiteId !== localSiteId
+			);
+			pruneEmptyConnectionsForUser( config, userKey );
+		}
+		await saveSharedConfig( config );
+	} finally {
+		await unlockSharedConfig();
+	}
+}
+
 /**
  * Updates specific connection entries in place (matched by remote site id and
  * local site id) for the current user. Entries that don't match an existing

packages/common/lib/tests/connected-sites.test.ts

@@ -0,0 +1,67 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { removeAllConnectedWpcomSitesForLocalSite } from '../connected-sites';
+import {
+	lockSharedConfig,
+	readSharedConfig,
+	saveSharedConfig,
+	unlockSharedConfig,
+} from '../shared-config';
+import type { SyncSite } from '../../types/sync';
+
+vi.mock( '../shared-config', () => ( {
+	lockSharedConfig: vi.fn(),
+	readSharedConfig: vi.fn(),
+	saveSharedConfig: vi.fn(),
+	unlockSharedConfig: vi.fn(),
+} ) );
+
+function connection( id: number, localSiteId: string ): SyncSite {
+	return {
+		id,
+		localSiteId,
+		name: `Site ${ id }`,
+		url: `https://site-${ id }.example.com`,
+		isStaging: false,
+		isPressable: false,
+		syncSupport: 'already-connected',
+		lastPullTimestamp: null,
+		lastPushTimestamp: null,
+	};
+}
+
+describe( 'removeAllConnectedWpcomSitesForLocalSite', () => {
+	beforeEach( () => {
+		vi.clearAllMocks();
+	} );
+
+	it( 'removes every connection for the local site across account buckets while locked', async () => {
+		const config = {
+			version: 1 as const,
+			connectedWpcomSites: {
+				'7': [ connection( 1, 'local-a' ), connection( 2, 'local-b' ), connection( 3, 'local-a' ) ],
+				'8': [ connection( 4, 'local-a' ) ],
+			},
+		};
+		vi.mocked( readSharedConfig ).mockResolvedValue( config );
+
+		await removeAllConnectedWpcomSitesForLocalSite( 'local-a' );
+
+		expect( lockSharedConfig ).toHaveBeenCalledOnce();
+		expect( saveSharedConfig ).toHaveBeenCalledWith( config );
+		expect( unlockSharedConfig ).toHaveBeenCalledOnce();
+		expect( config.connectedWpcomSites[ '7' ] ).toEqual( [ connection( 2, 'local-b' ) ] );
+		expect( config.connectedWpcomSites ).not.toHaveProperty( '8' );
+	} );
+
+	it( 'prunes the connection map when no connections remain', async () => {
+		const config = {
+			version: 1 as const,
+			connectedWpcomSites: { '7': [ connection( 1, 'local-a' ) ] },
+		};
+		vi.mocked( readSharedConfig ).mockResolvedValue( config );
+
+		await removeAllConnectedWpcomSitesForLocalSite( 'local-a' );
+
+		expect( config ).not.toHaveProperty( 'connectedWpcomSites' );
+	} );
+} );