From 9f1d0889bc9d70b6540200340b1ecb2497cddddd Mon Sep 17 00:00:00 2001
From: Grzegorz Zdunek <grzegorz.zdunek@goteleport.com>
Date: Fri, 18 Mar 2022 10:48:51 +0100
Subject: [PATCH] Take `localClusterUri` into account in `QuickInput`

---
 .../src/ui/QuickInput/useQuickInput.ts         |  5 +++--
 .../quickInput/quickInputService.test.ts       | 14 ++++++++++++++
 .../src/ui/services/quickInput/quickPickers.ts | 18 ++++++++++++------
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/packages/teleterm/src/ui/QuickInput/useQuickInput.ts b/packages/teleterm/src/ui/QuickInput/useQuickInput.ts
index 5a3319678..5f57d232b 100644
--- a/packages/teleterm/src/ui/QuickInput/useQuickInput.ts
+++ b/packages/teleterm/src/ui/QuickInput/useQuickInput.ts
@@ -26,7 +26,6 @@ export default function useQuickInput() {
   const {
     quickInputService,
     workspacesService,
-    clustersService,
     commandLauncher,
   } = useAppContext();
   workspacesService.useState();
@@ -36,7 +35,9 @@ export default function useQuickInput() {
   const [activeSuggestion, setActiveSuggestion] = React.useState(0);
   const autocompleteResult = React.useMemo(
     () => quickInputService.getAutocompleteResult(inputValue),
-    [inputValue]
+    // `localClusterUri` has been added to refresh suggestions from
+    // `QuickSshLoginPicker` and `QuickServerPicker` when it changes
+    [inputValue, workspacesService.getActiveWorkspace()?.localClusterUri]
   );
   const hasSuggestions =
     autocompleteResult.kind === 'autocomplete.partial-match';
diff --git a/packages/teleterm/src/ui/services/quickInput/quickInputService.test.ts b/packages/teleterm/src/ui/services/quickInput/quickInputService.test.ts
index c47d99d59..bf978cdbe 100644
--- a/packages/teleterm/src/ui/services/quickInput/quickInputService.test.ts
+++ b/packages/teleterm/src/ui/services/quickInput/quickInputService.test.ts
@@ -299,6 +299,13 @@ test('getAutocompleteResult returns correct result for an SSH host suggestion ri
     CommandLauncherMock,
     onlyTshSshCommand
   );
+  jest
+    .spyOn(WorkspacesServiceMock.prototype, 'getActiveWorkspace')
+    .mockImplementation(() => ({
+      localClusterUri: 'test_uri',
+      documents: [],
+      location: '',
+    }));
   jest
     .spyOn(ClustersServiceMock.prototype, 'searchServers')
     .mockImplementation(() => {
@@ -351,6 +358,13 @@ test('getAutocompleteResult returns correct result for a partial match on an SSH
         },
       ];
     });
+  jest
+    .spyOn(WorkspacesServiceMock.prototype, 'getActiveWorkspace')
+    .mockImplementation(() => ({
+      localClusterUri: 'test_uri',
+      documents: [],
+      location: '',
+    }));
   const quickInputService = new QuickInputService(
     new CommandLauncherMock(undefined),
     new ClustersServiceMock(undefined),
diff --git a/packages/teleterm/src/ui/services/quickInput/quickPickers.ts b/packages/teleterm/src/ui/services/quickInput/quickPickers.ts
index 00204b8b1..892d5af82 100644
--- a/packages/teleterm/src/ui/services/quickInput/quickPickers.ts
+++ b/packages/teleterm/src/ui/services/quickInput/quickPickers.ts
@@ -257,10 +257,13 @@ export class QuickSshLoginPicker implements QuickInputPicker {
   ) {}
 
   private filterSshLogins(input: string): SuggestionSshLogin[] {
-    // TODO(ravicious): Use local cluster URI.
     // TODO(ravicious): Handle the `--cluster` tsh ssh flag.
-    const rootClusterUri = this.workspacesService.getRootClusterUri();
-    const cluster = this.clustersService.findCluster(rootClusterUri);
+    const localClusterUri =
+      this.workspacesService.getActiveWorkspace()?.localClusterUri;
+    if (!localClusterUri) {
+      return [];
+    }
+    const cluster = this.clustersService.findCluster(localClusterUri);
     const allLogins = cluster?.loggedInUser?.sshLoginsList || [];
     let matchingLogins: typeof allLogins;
 
@@ -302,10 +305,13 @@ export class QuickServerPicker implements QuickInputPicker {
   ) {}
 
   private filterServers(input: string): SuggestionServer[] {
-    // TODO(ravicious): Use local cluster URI.
     // TODO(ravicious): Handle the `--cluster` tsh ssh flag.
-    const rootClusterUri = this.workspacesService.getRootClusterUri();
-    const servers = this.clustersService.searchServers(rootClusterUri, {
+    const localClusterUri =
+      this.workspacesService.getActiveWorkspace()?.localClusterUri;
+    if (!localClusterUri) {
+      return [];
+    }
+    const servers = this.clustersService.searchServers(localClusterUri, {
       search: input,
     });