Combine ClustersService logout functions

web/packages/teleterm/src/ui/ClusterLogout/ClusterLogout.tsx

@@ -25,10 +25,12 @@ import DialogConfirmation, {
 } from 'design/DialogConfirmation';
 import { Cross } from 'design/Icon';
 import { P } from 'design/Text/Text';
+import { useAsync } from 'shared/hooks/useAsync';
 
+import { useAppContext } from 'teleterm/ui/appContextProvider';
 import { RootClusterUri } from 'teleterm/ui/uri';
 
-import { useClusterLogout } from './useClusterLogout';
+import { logoutWithCleanup } from './logoutWithCleanup';
 
 export function ClusterLogout({
   clusterUri,
@@ -41,9 +43,10 @@ export function ClusterLogout({
   hidden?: boolean;
   onClose(): void;
 }) {
-  const { removeCluster, status, statusText } = useClusterLogout({
-    clusterUri,
-  });
+  const ctx = useAppContext();
+  const [{ status, statusText }, removeCluster] = useAsync(() =>
+    logoutWithCleanup(ctx, clusterUri)
+  );
 
   async function removeClusterAndClose(): Promise<void> {
     const [, err] = await removeCluster();

web/packages/teleterm/src/ui/ClusterLogout/logoutWithCleanup.ts

@@ -0,0 +1,69 @@
+/**
+ * Teleport
+ * Copyright (C) 2025 Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+import Logger from 'teleterm/logger';
+import { IAppContext } from 'teleterm/ui/types';
+import { RootClusterUri, routing } from 'teleterm/ui/uri';
+
+/** Disposes cluster-related resources and then logs out. */
+export async function logoutWithCleanup(
+  ctx: IAppContext,
+  clusterUri: RootClusterUri
+): Promise<void> {
+  const logger = new Logger('logoutWithCleanup');
+  // This function checks for updates, do not wait for it.
+  ctx.mainProcessClient
+    .maybeRemoveAppUpdatesManagingCluster(clusterUri)
+    .catch(err => {
+      logger.error('Failed to remove managing cluster', err);
+    });
+
+  if (ctx.workspacesService.getRootClusterUri() === clusterUri) {
+    const [firstConnectedWorkspace] = ctx.workspacesService
+      .getConnectedWorkspacesClustersUri()
+      .filter(c => c !== clusterUri);
+    if (firstConnectedWorkspace) {
+      await ctx.workspacesService.setActiveWorkspace(firstConnectedWorkspace);
+    } else {
+      await ctx.workspacesService.setActiveWorkspace(null);
+    }
+  }
+
+  // Remove connections first, they depend both on the cluster and the workspace.
+  ctx.connectionTracker.removeItemsBelongingToRootCluster(clusterUri);
+  // Remove the workspace next, because it depends on the cluster.
+  ctx.workspacesService.removeWorkspace(clusterUri);
+
+  // If there are active ssh connections to the agent, killing it will take a few seconds. To work
+  // around this, kill the agent only after removing the workspace. Removing the workspace closes
+  // ssh tabs, so it should terminate connections to the cluster from the app.
+  await ctx.connectMyComputerService.killAgentAndRemoveData(clusterUri);
+
+  await ctx.clustersService.removeClusterGateways(clusterUri);
+
+  const {
+    params: { rootClusterId },
+  } = routing.parseClusterUri(clusterUri);
+  await ctx.mainProcessClient.removeKubeConfig({
+    relativePath: rootClusterId,
+    isDirectory: true,
+  });
+
+  // Remove the cluster, it does not depend on anything.
+  await ctx.clustersService.logout(clusterUri);
+}

web/packages/teleterm/src/ui/services/clusters/clustersService.test.ts

@@ -136,7 +136,7 @@ test('add cluster does not overwrite the existing cluster', async () => {
   );
 });
 
-test('remove cluster', async () => {
+test('remove gateways', async () => {
   const { removeGateway } = getClientMocks();
   const service = createService({ removeGateway });
   const gatewayFromRootCluster = makeDatabaseGateway({
@@ -164,14 +164,11 @@ test('remove cluster', async () => {
     ]);
   });
 
-  await service.removeClusterAndResources(clusterUri);
+  await service.removeClusterGateways(clusterUri);
 
-  expect(service.findCluster(clusterUri)).toBeUndefined();
-  expect(service.findCluster(leafClusterMock.uri)).toBeUndefined();
   expect(service.state.gateways).toEqual(
     new Map([[gatewayFromOtherCluster.uri, gatewayFromOtherCluster]])
   );
-
   expect(removeGateway).toHaveBeenCalledWith({
     gatewayUri: gatewayFromRootCluster.uri,
   });
@@ -222,8 +219,8 @@ test('logout from cluster', async () => {
 
   expect(logout).toHaveBeenCalledWith({ clusterUri });
   expect(removeCluster).toHaveBeenCalledWith({ clusterUri });
-  expect(service.findCluster(clusterMock.uri).connected).toBe(false);
-  expect(service.findCluster(leafClusterMock.uri).connected).toBe(false);
+  expect(service.findCluster(clusterMock.uri)).toBeUndefined();
+  expect(service.findCluster(leafClusterMock.uri)).toBeUndefined();
 });
 
 test('create a gateway', async () => {

web/packages/teleterm/src/ui/services/clusters/clustersService.ts

@@ -80,15 +80,7 @@ export class ClustersService extends ImmutableStore<ClustersServiceState> {
     return cluster;
   }
 
-  /**
-   * Logs out of the cluster and removes the profile.
-   * Does not remove the cluster from the state, but sets the cluster and its leafs as disconnected.
-   * It needs to be done, because some code can operate on the cluster the intermediate period between logout
-   * and actually removing it from the state.
-   * A code that operates on that intermediate state is in `useClusterLogout.tsx`.
-   * After invoking `logout()`, it looks for the next workspace to switch to. If we hadn't marked the cluster as disconnected,
-   * the method might have returned us the same cluster we wanted to log out of.
-   */
+  /** Logs out of the cluster. */
   async logout(clusterUri: uri.RootClusterUri) {
     // TODO(gzdunek): logout and removeCluster should be combined into a single acton in tshd
     await this.client.logout({ clusterUri });
@@ -97,7 +89,7 @@ export class ClustersService extends ImmutableStore<ClustersServiceState> {
     this.setState(draft => {
       draft.clusters.forEach(cluster => {
         if (routing.belongsToProfile(clusterUri, cluster.uri)) {
-          cluster.connected = false;
+          draft.clusters.delete(cluster.uri);
         }
       });
     });
@@ -192,6 +184,13 @@ export class ClustersService extends ImmutableStore<ClustersServiceState> {
     ]);
   }
 
+  /**
+   * Synchronizes root clusters.
+   *
+   * This should only be called before creating workspaces.
+   * If called afterward, a cluster might be removed without first removing
+   * its associated workspace, resulting in an invalid state.
+   */
   async syncRootClustersAndCatchErrors(abortSignal?: AbortSignal) {
     let clusters: Cluster[];
 
@@ -315,22 +314,9 @@ export class ClustersService extends ImmutableStore<ClustersServiceState> {
     return response;
   }
 
-  /** Removes cluster, its leafs and other resources. */
-  async removeClusterAndResources(clusterUri: uri.RootClusterUri) {
-    this.setState(draft => {
-      draft.clusters.forEach(cluster => {
-        if (routing.belongsToProfile(clusterUri, cluster.uri)) {
-          draft.clusters.delete(cluster.uri);
-        }
-      });
-    });
-    await this.removeClusterKubeConfigs(clusterUri);
-    await this.removeClusterGateways(clusterUri);
-  }
-
   // TODO(ravicious): Create a single RPC for this rather than sending a separate request for each
   // gateway.
-  private async removeClusterGateways(clusterUri: uri.RootClusterUri) {
+  async removeClusterGateways(clusterUri: uri.RootClusterUri) {
     for (const [, gateway] of this.state.gateways) {
       if (routing.belongsToProfile(clusterUri, gateway.targetUri)) {
         try {
@@ -512,16 +498,6 @@ export class ClustersService extends ImmutableStore<ClustersServiceState> {
     return this.getClusters().filter(c => !c.leaf);
   }
 
-  async removeClusterKubeConfigs(clusterUri: string): Promise<void> {
-    const {
-      params: { rootClusterId },
-    } = routing.parseClusterUri(clusterUri);
-    return this.mainProcessClient.removeKubeConfig({
-      relativePath: rootClusterId,
-      isDirectory: true,
-    });
-  }
-
   useState() {
     return useStore(this).state;
   }

web/packages/teleterm/src/ui/services/workspacesService/workspacesService.ts

@@ -352,6 +352,7 @@ export class WorkspacesService extends ImmutableStore<WorkspacesState> {
     }
 
     if (cluster.profileStatusError) {
+      // TODO(gzdunek): We should only sync the target cluster, not all of them.
       await this.clustersService.syncRootClustersAndCatchErrors(abortSignal);
       // Update the cluster.
       cluster = this.clustersService.findCluster(clusterUri);