Connect: Fix logout sequence (#24978)

* Fix logout sequence

* Adjust `useLoggedInUser` documentation

* Mark clusters as disconnected after logging out

* Refactor `ClusterLogout` to not use the hook/container pattern

* Run prettier

* Fix test

* Use `routing.belongsToProfile`

web/packages/teleterm/src/services/tshd/fixtures/mocks.ts

@@ -81,7 +81,7 @@ export class MockTshClient implements TshClient {
 
   getCluster: (clusterUri: string) => Promise<Cluster>;
   getAuthSettings: (clusterUri: string) => Promise<AuthSettings>;
-  removeCluster: (clusterUri: string) => Promise<undefined>;
+  removeCluster = () => Promise.resolve();
   loginLocal: (
     params: LoginLocalParams,
     abortSignal?: TshAbortSignal
@@ -94,7 +94,7 @@ export class MockTshClient implements TshClient {
     params: LoginPasswordlessParams,
     abortSignal?: TshAbortSignal
   ) => Promise<undefined>;
-  logout: (clusterUri: string) => Promise<undefined>;
+  logout = () => Promise.resolve();
   transferFile: () => undefined;
   reportUsageEvent: () => undefined;
 }

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

@@ -16,6 +16,8 @@
 
 import React from 'react';
 
+import { MockAppContextProvider } from 'teleterm/ui/fixtures/MockAppContextProvider';
+
 import { ClusterLogout } from './ClusterLogout';
 
 export default {
@@ -24,13 +26,12 @@ export default {
 
 export const Story = () => {
   return (
-    <ClusterLogout
-      status=""
-      statusText=""
-      removeCluster={() => Promise.resolve([undefined, null])}
-      onClose={() => {}}
-      clusterUri="/clusters/foo"
-      clusterTitle="foo"
-    />
+    <MockAppContextProvider>
+      <ClusterLogout
+        onClose={() => {}}
+        clusterUri="/clusters/foo"
+        clusterTitle="foo"
+      />
+    </MockAppContextProvider>
   );
 };

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

@@ -25,20 +25,32 @@ import { ButtonIcon, ButtonPrimary, Text } from 'design';
 
 import { Close } from 'design/Icon';
 
-import { Props, State, useClusterLogout } from './useClusterLogout';
+import { RootClusterUri } from 'teleterm/ui/uri';
 
-export default function Container(props: Props) {
-  const state = useClusterLogout(props);
-  return <ClusterLogout {...state} />;
+import { useClusterLogout } from './useClusterLogout';
+
+interface ClusterLogoutProps {
+  clusterTitle: string;
+  clusterUri: RootClusterUri;
+  onClose(): void;
 }
 
 export function ClusterLogout({
-  status,
+  clusterUri,
   onClose,
-  statusText,
   clusterTitle,
-  removeCluster,
-}: State) {
+}: ClusterLogoutProps) {
+  const { removeCluster, status, statusText } = useClusterLogout({
+    clusterUri,
+  });
+
+  async function removeClusterAndClose(): Promise<void> {
+    const [, err] = await removeCluster();
+    if (!err) {
+      onClose();
+    }
+  }
+
   return (
     <DialogConfirmation
       open={true}
@@ -51,7 +63,7 @@ export function ClusterLogout({
       <form
         onSubmit={e => {
           e.preventDefault();
-          removeCluster();
+          removeClusterAndClose();
         }}
       >
         <DialogHeader justifyContent="space-between" mb={0}>

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

@@ -15,13 +15,16 @@
  */
 
 import { useAsync } from 'shared/hooks/useAsync';
-import { useEffect } from 'react';
 
 import { RootClusterUri } from 'teleterm/ui/uri';
 
 import { useAppContext } from '../appContextProvider';
 
-export function useClusterLogout({ clusterUri, onClose, clusterTitle }: Props) {
+export function useClusterLogout({
+  clusterUri,
+}: {
+  clusterUri: RootClusterUri;
+}) {
   const ctx = useAppContext();
   const [{ status, statusText }, removeCluster] = useAsync(async () => {
     await ctx.clustersService.logout(clusterUri);
@@ -35,30 +38,18 @@ export function useClusterLogout({ clusterUri, onClose, clusterTitle }: Props) {
         await ctx.workspacesService.setActiveWorkspace(null);
       }
     }
-    ctx.workspacesService.removeWorkspace(clusterUri);
-    ctx.connectionTracker.removeItemsBelongingToRootCluster(clusterUri);
-  });
 
-  useEffect(() => {
-    if (status === 'success') {
-      onClose();
-    }
-  }, [status]);
+    // 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);
+    // remove the cluster, it does not depend on anything
+    await ctx.clustersService.removeClusterAndResources(clusterUri);
+  });
 
   return {
     status,
     statusText,
     removeCluster,
-    onClose,
-    clusterUri,
-    clusterTitle,
   };
 }
-
-export type Props = {
-  onClose?(): void;
-  clusterTitle?: string;
-  clusterUri: RootClusterUri;
-};
-
-export type State = ReturnType<typeof useClusterLogout>;

web/packages/teleterm/src/ui/ModalsHost/ModalsHost.test.tsx

@@ -39,15 +39,14 @@ const documentsReopenDialog: DialogDocumentsReopen = {
   onCancel: () => {},
 };
 
-jest.mock('teleterm/ui/ClusterLogout/ClusterLogout', () => {
-  const MockClusterLogout = () => (
+jest.mock('teleterm/ui/ClusterLogout', () => ({
+  ClusterLogout: () => (
     <div
       data-testid="mocked-dialog"
       data-dialog-kind={clusterLogoutDialog.kind}
     ></div>
-  );
-  return MockClusterLogout;
-});
+  ),
+}));
 
 jest.mock('teleterm/ui/DocumentsReopen', () => ({
   DocumentsReopen: () => (

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

@@ -22,7 +22,7 @@ import { ClusterConnect } from 'teleterm/ui/ClusterConnect';
 import { DocumentsReopen } from 'teleterm/ui/DocumentsReopen';
 import { Dialog } from 'teleterm/ui/services/modals';
 
-import ClusterLogout from '../ClusterLogout/ClusterLogout';
+import { ClusterLogout } from '../ClusterLogout';
 import { ResourceSearchErrors } from '../Search/ResourceSearchErrors';
 import { assertUnreachable } from '../utils';
 

web/packages/teleterm/src/ui/hooks/useLoggedInUser.ts

@@ -25,8 +25,7 @@ import { LoggedInUser } from 'teleterm/services/tshd/types';
  * It should be used within components that reside outside of WorkspaceContext, typically anything
  * that's outside of Document-type components.
  *
- * It might return undefined if there's no active workspace or during the logout procedure because
- * ClustersService state is cleared up before WorkspacesService state.
+ * It might return undefined if there's no active workspace.
  */
 export function useLoggedInUser(): LoggedInUser | undefined {
   const { clustersService, workspacesService } = useAppContext();
@@ -51,10 +50,8 @@ export function useLoggedInUser(): LoggedInUser | undefined {
  * rendered inside of them.
  *
  * In general, the callsite should always assume that this function might return undefined.
- * One case where it will for sure return undefined is during the logout process as
- * ClustersService state is cleared up before WorkspacesService state. On top of that, each
- * workspace is always rendered, even when the cluster is not connected, with at least the default
- * document. In that scenario useWorkspaceLoggedInUser could return undefined when used within the
+ * Each workspace is always rendered, even when the cluster is not connected, with at least the default
+ * document. In that scenario, useWorkspaceLoggedInUser could return undefined when used within the
  * default document.
  */
 export function useWorkspaceLoggedInUser(): LoggedInUser | undefined {

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

@@ -123,10 +123,7 @@ test('add cluster', async () => {
 });
 
 test('remove cluster', async () => {
-  const { removeCluster } = getClientMocks();
-  const service = createService({
-    removeCluster,
-  });
+  const service = createService({});
 
   service.setState(draftState => {
     draftState.clusters = new Map([
@@ -135,9 +132,8 @@ test('remove cluster', async () => {
     ]);
   });
 
-  await service.removeCluster(clusterUri);
+  await service.removeClusterAndResources(clusterUri);
 
-  expect(removeCluster).toHaveBeenCalledWith(clusterUri);
   expect(service.findCluster(clusterUri)).toBeUndefined();
   expect(service.findCluster(leafClusterMock.uri)).toBeUndefined();
 });
@@ -192,8 +188,9 @@ test('logout from cluster', async () => {
   await service.logout(clusterUri);
 
   expect(logout).toHaveBeenCalledWith(clusterUri);
-  expect(service.findCluster(clusterUri)).toBeUndefined();
-  expect(service.findCluster(leafClusterMock.uri)).toBeUndefined();
+  expect(removeCluster).toHaveBeenCalledWith(clusterUri);
+  expect(service.findCluster(clusterMock.uri).connected).toBe(false);
+  expect(service.findCluster(leafClusterMock.uri).connected).toBe(false);
 });
 
 test('create a gateway', async () => {

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

@@ -71,11 +71,27 @@ export class ClustersService extends ImmutableStore<types.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.
+   */
   async logout(clusterUri: uri.RootClusterUri) {
     // TODO(gzdunek): logout and removeCluster should be combined into a single acton in tshd
     await this.client.logout(clusterUri);
-    await this.removeCluster(clusterUri);
-    await this.removeClusterKubeConfigs(clusterUri);
+    await this.client.removeCluster(clusterUri);
+
+    this.setState(draft => {
+      draft.clusters.forEach(cluster => {
+        if (routing.belongsToProfile(clusterUri, cluster.uri)) {
+          cluster.connected = false;
+        }
+      });
+    });
   }
 
   async loginLocal(
@@ -301,23 +317,16 @@ export class ClustersService extends ImmutableStore<types.ClustersServiceState>
     return response;
   }
 
-  /**
-   * Removes cluster and its leaf clusters (if any)
-   */
-  async removeCluster(clusterUri: uri.RootClusterUri) {
-    await this.client.removeCluster(clusterUri);
-    const leafClustersUris = this.getClusters()
-      .filter(
-        item =>
-          item.leaf && routing.ensureRootClusterUri(item.uri) === clusterUri
-      )
-      .map(cluster => cluster.uri);
+  /** Removes cluster, its leafs and other resources. */
+  async removeClusterAndResources(clusterUri: uri.RootClusterUri) {
     this.setState(draft => {
-      draft.clusters.delete(clusterUri);
-      leafClustersUris.forEach(leafClusterUri => {
-        draft.clusters.delete(leafClusterUri);
+      draft.clusters.forEach(cluster => {
+        if (routing.belongsToProfile(clusterUri, cluster.uri)) {
+          draft.clusters.delete(cluster.uri);
+        }
       });
     });
+    await this.removeClusterKubeConfigs(clusterUri);
   }
 
   async getAuthSettings(clusterUri: uri.RootClusterUri) {