Remove state related to a cluster when removing it (#755)

packages/teleterm/src/ui/appContext.ts

@@ -55,7 +55,7 @@ export default class AppContext {
       this.modalsService,
       this.clustersService,
       this.notificationsService,
-      this.statePersistenceService,
+      this.statePersistenceService
     );
     this.terminalsService = new TerminalsService(ptyServiceClient);
 
@@ -81,9 +81,6 @@ export default class AppContext {
 
   async init(): Promise<void> {
     await this.clustersService.syncRootClusters();
-    const { rootClusterUri } = this.statePersistenceService.getWorkspaces();
-    if (rootClusterUri) {
-      this.workspacesService.setActiveWorkspace(rootClusterUri);
-    }
+    this.workspacesService.restorePersistedState();
   }
 }

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

@@ -278,12 +278,28 @@ export class ClustersService extends ImmutableStore<ClustersServiceState> {
     }
   }
 
+  /**
+   * Removes cluster and its leaf clusters (if any)
+   */
   async removeCluster(clusterUri: string) {
     await this.client.removeCluster(clusterUri);
+    const leafClustersUris = this.getClusters()
+      .filter(
+        item =>
+          item.leaf && routing.ensureRootClusterUri(item.uri) === clusterUri
+      )
+      .map(cluster => cluster.uri);
     this.setState(draft => {
       draft.clusters.delete(clusterUri);
+      leafClustersUris.forEach(leafClusterUri => {
+        draft.clusters.delete(leafClusterUri);
+      });
     });
+
     this.removeResources(clusterUri);
+    leafClustersUris.forEach(leafClusterUri => {
+      this.removeResources(leafClusterUri);
+    });
   }
 
   async getAuthSettings(clusterUri: string) {

packages/teleterm/src/ui/services/connectionTracker/connectionTrackerService.ts

@@ -102,6 +102,12 @@ export class ConnectionTrackerService extends ImmutableStore<ConnectionTrackerSt
 
   private _refreshState = () => {
     this.setState(draft => {
+      // filter out connections from removed clusters
+      draft.connections = draft.connections.filter(i => {
+        const uri = i.kind === 'connection.gateway' ? i.targetUri : i.serverUri;
+        return !!this._clusterService.findClusterByResource(uri);
+      });
+
       // assign default "connected" values
       draft.connections.forEach(i => {
         if (i.kind === 'connection.gateway') {
@@ -134,6 +140,9 @@ export class ConnectionTrackerService extends ImmutableStore<ConnectionTrackerSt
         switch (doc.kind) {
           // process gateway connections
           case 'doc.gateway':
+            if (!doc.port) {
+              break;
+            }
             const gwConn = draft.connections.find(
               getGatewayConnectionByDocument(doc)
             ) as TrackedGatewayConnection;

packages/teleterm/src/ui/services/statePersistence/statePersistenceService.ts

@@ -20,50 +20,49 @@ import { WorkspacesState } from 'teleterm/ui/services/workspacesService';
 
 interface StatePersistenceState {
   connectionTracker: ConnectionTrackerState;
-  workspacesState?: WorkspacesState;
+  workspacesState: WorkspacesState;
 }
 
 export class StatePersistenceService {
-  state: StatePersistenceState = {
-    connectionTracker: {
-      connections: [],
-    },
-    workspacesState: {
-      workspaces: {},
-    },
-  };
+  constructor(private _fileStorage: FileStorage) {}
 
-  constructor(private _fileStorage: FileStorage) {
-    const restored = this._fileStorage.get<StatePersistenceState>('state');
-    if (restored) {
-      this.state = restored;
-    }
+  saveConnectionTrackerState(connectionTracker: ConnectionTrackerState): void {
+    const newState: StatePersistenceState = {
+      ...this.getState(),
+      connectionTracker,
+    };
+    this.putState(newState);
   }
 
-  saveConnectionTrackerState(navigatorState: ConnectionTrackerState): void {
-    this.state.connectionTracker = navigatorState;
-    this._fileStorage.put('state', this.state);
+  getConnectionTrackerState(): ConnectionTrackerState {
+    return this.getState().connectionTracker;
   }
 
-  getConnectionTrackerState(): ConnectionTrackerState {
-    return this.state.connectionTracker;
+  saveWorkspacesState(workspacesState: WorkspacesState): void {
+    const newState: StatePersistenceState = {
+      ...this.getState(),
+      workspacesState,
+    };
+    this.putState(newState);
+  }
+
+  getWorkspacesState(): WorkspacesState {
+    return this.getState().workspacesState;
   }
 
-  saveWorkspaces(workspacesState: WorkspacesState): void {
-    this.state.workspacesState.rootClusterUri = workspacesState.rootClusterUri;
-    for (let w in workspacesState.workspaces) {
-      if (workspacesState.workspaces[w]) {
-        this.state.workspacesState.workspaces[w] = {
-          location: workspacesState.workspaces[w].location,
-          localClusterUri: workspacesState.workspaces[w].localClusterUri,
-          documents: workspacesState.workspaces[w].documents,
-        };
-      }
-    }
-    this._fileStorage.put('state', this.state);
+  private getState(): StatePersistenceState {
+    const defaultState: StatePersistenceState = {
+      connectionTracker: {
+        connections: [],
+      },
+      workspacesState: {
+        workspaces: {},
+      },
+    };
+    return this._fileStorage.get('state') || defaultState;
   }
 
-  getWorkspaces(): WorkspacesState {
-    return this.state.workspacesState;
+  private putState(state: StatePersistenceState): void {
+    this._fileStorage.put('state', state);
   }
 }

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

@@ -89,14 +89,16 @@ export class DocumentsService {
   }
 
   createGatewayDocument(opts: CreateGatewayDocumentOpts): DocumentGateway {
-    const { targetUri, title, targetUser } = opts;
+    const { targetUri, title, targetUser, port, gatewayUri } = opts;
     const uri = routing.getDocUri({ docId: unique() });
     return {
       uri,
       kind: 'doc.gateway',
       targetUri,
       targetUser,
+      gatewayUri,
       title,
+      port
     };
   }
 

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

@@ -6,6 +6,7 @@ import { ClustersService } from 'teleterm/ui/services/clusters';
 import { StatePersistenceService } from 'teleterm/ui/services/statePersistence';
 import { isEqual } from 'lodash';
 import { NotificationsService } from 'teleterm/ui/services/notifications';
+import { routing } from 'teleterm/ui/uri';
 
 export interface WorkspacesState {
   rootClusterUri?: string;
@@ -107,30 +108,16 @@ export class WorkspacesService extends ImmutableStore<WorkspacesState> {
 
   setState(nextState: (draftState: WorkspacesState) => WorkspacesState | void) {
     super.setState(nextState);
-    this.statePersistenceService.saveWorkspaces(this.state);
+    this.persistState();
   }
 
   setActiveWorkspace(clusterUri: string): Promise<void> {
     const setWorkspace = () => {
       this.setState(draftState => {
+        // adding a new workspace
         if (!draftState.workspaces[clusterUri]) {
-          const persistedWorkspace =
-            this.statePersistenceService.getWorkspaces().workspaces[clusterUri];
-          const defaultDocument = this.getWorkspaceDocumentService(
-            clusterUri
-          ).createClusterDocument({ clusterUri });
-
-          draftState.workspaces[clusterUri] = {
-            localClusterUri: persistedWorkspace?.localClusterUri || clusterUri,
-            location: defaultDocument.uri,
-            documents: [defaultDocument],
-            previous: persistedWorkspace?.documents
-              ? {
-                  documents: persistedWorkspace.documents,
-                  location: persistedWorkspace.location,
-                }
-              : undefined,
-          };
+          draftState.workspaces[clusterUri] =
+            this.getWorkspaceDefaultState(clusterUri);
         }
         draftState.rootClusterUri = clusterUri;
       });
@@ -174,8 +161,7 @@ export class WorkspacesService extends ImmutableStore<WorkspacesState> {
     })
       .then(() => {
         return new Promise<void>(resolve => {
-          if (!this.canReopenPreviousDocuments(this.getWorkspace(clusterUri))) {
-            this.discardPreviousDocuments(clusterUri);
+          if (!this.getWorkspace(clusterUri)?.previous) {
             return resolve();
           }
           this.modalsService.openDocumentsReopenDialog({
@@ -205,6 +191,41 @@ export class WorkspacesService extends ImmutableStore<WorkspacesState> {
     );
   }
 
+  restorePersistedState(): void {
+    const persistedState = this.statePersistenceService.getWorkspacesState();
+    const restoredWorkspaces = this.clustersService
+      .getClusters()
+      .reduce((workspaces, cluster) => {
+        const persistedWorkspace = persistedState.workspaces[cluster.uri];
+        const workspaceDefaultState = this.getWorkspaceDefaultState(
+          persistedWorkspace?.localClusterUri || cluster.uri
+        );
+        const persistedWorkspaceDocuments = persistedWorkspace.documents;
+
+        workspaces[cluster.uri] = {
+          ...workspaceDefaultState,
+          previous: this.canReopenPreviousDocuments({
+            previousDocuments: persistedWorkspaceDocuments,
+            currentDocuments: workspaceDefaultState.documents,
+          })
+            ? {
+                location: persistedWorkspace.location,
+                documents: persistedWorkspaceDocuments,
+              }
+            : undefined,
+        };
+        return workspaces;
+      }, {});
+
+    this.setState(draftState => {
+      draftState.workspaces = restoredWorkspaces;
+    });
+
+    if (persistedState.rootClusterUri) {
+      this.setActiveWorkspace(persistedState.rootClusterUri);
+    }
+  }
+
   private reopenPreviousDocuments(clusterUri: string): void {
     this.setState(draftState => {
       const workspace = draftState.workspaces[clusterUri];
@@ -221,17 +242,47 @@ export class WorkspacesService extends ImmutableStore<WorkspacesState> {
     });
   }
 
-  private canReopenPreviousDocuments(workspace: Workspace): boolean {
+  private canReopenPreviousDocuments({
+    previousDocuments,
+    currentDocuments,
+  }: {
+    previousDocuments: Document[];
+    currentDocuments: Document[];
+  }): boolean {
     const removeUri = (documents: Document[]) =>
       documents.map(d => ({ ...d, uri: undefined }));
 
     return (
-      workspace.previous &&
-      workspace.previous.documents?.length &&
-      !isEqual(
-        removeUri(workspace.previous.documents),
-        removeUri(workspace.documents)
-      )
+      previousDocuments?.length &&
+      !isEqual(removeUri(previousDocuments), removeUri(currentDocuments))
     );
   }
+
+  private getWorkspaceDefaultState(localClusterUri: string): Workspace {
+    const rootClusterUri = routing.ensureRootClusterUri(localClusterUri);
+    const defaultDocument = this.getWorkspaceDocumentService(
+      rootClusterUri
+    ).createClusterDocument({ clusterUri: localClusterUri });
+    return {
+      localClusterUri,
+      location: defaultDocument.uri,
+      documents: [defaultDocument],
+    };
+  }
+
+  private persistState(): void {
+    const stateToSave: WorkspacesState = {
+      rootClusterUri: this.state.rootClusterUri,
+      workspaces: {},
+    };
+    for (let w in this.state.workspaces) {
+      const workspace = this.state.workspaces[w];
+      stateToSave.workspaces[w] = {
+        localClusterUri: workspace.localClusterUri,
+        location: workspace.previous?.location || workspace.location,
+        documents: workspace.previous?.documents || workspace.documents,
+      };
+    }
+    this.statePersistenceService.saveWorkspacesState(stateToSave);
+  }
 }