Connect: Replace legacy useStore hook with useStoreSelector in Worksp…

…acesService (#48353)

* Add useWorkspaceServiceState hook

* Do trivial replacements of WorkspacesService.useState

* Replace useState with useStoreSelector and actually select some state

* Refactor useWorkspaceContext to use useStoreSelector

* Remove WorkspaceContextProvider from connectMyComputerContext tests

Workspace context is not used by connect my computer context.

* Remove useState from WorkspacesService

web/packages/teleterm/src/ui/AccessRequestCheckout/useAccessRequestCheckout.ts

@@ -33,6 +33,7 @@ import { useSpecifiableFields } from 'shared/components/AccessRequests/NewReques
 import { CreateRequest } from 'shared/components/AccessRequests/Shared/types';
 
 import { useAppContext } from 'teleterm/ui/appContextProvider';
+import { useWorkspaceServiceState } from 'teleterm/ui/services/workspacesService';
 import {
   PendingAccessRequest,
   extractResourceRequestProperties,
@@ -54,7 +55,7 @@ import { makeUiAccessRequest } from '../DocumentAccessRequests/useAccessRequests
 
 export default function useAccessRequestCheckout() {
   const ctx = useAppContext();
-  ctx.workspacesService.useState();
+  useWorkspaceServiceState();
   ctx.clustersService.useState();
   const clusterUri =
     ctx.workspacesService?.getActiveWorkspace()?.localClusterUri;

web/packages/teleterm/src/ui/ConnectMyComputer/connectMyComputerContext.test.tsx

@@ -18,13 +18,11 @@
 
 import { EventEmitter } from 'node:events';
 
-import React from 'react';
 import { act, renderHook, waitFor } from '@testing-library/react';
 import { makeErrorAttempt } from 'shared/hooks/useAsync';
 
 import { MockAppContextProvider } from 'teleterm/ui/fixtures/MockAppContextProvider';
 import { MockAppContext } from 'teleterm/ui/fixtures/mocks';
-import { WorkspaceContextProvider } from 'teleterm/ui/Documents';
 import { AgentProcessState } from 'teleterm/mainProcess/types';
 import * as resourcesContext from 'teleterm/ui/DocumentCluster/resourcesContext';
 import {
@@ -90,13 +88,11 @@ function renderUseConnectMyComputerContextHook(
   return renderHook(() => useConnectMyComputerContext(), {
     wrapper: ({ children }) => (
       <MockAppContextProvider appContext={appContext}>
-        <WorkspaceContextProvider value={null}>
-          <resourcesContext.ResourcesContextProvider>
-            <ConnectMyComputerContextProvider rootClusterUri={rootCluster.uri}>
-              {children}
-            </ConnectMyComputerContextProvider>
-          </resourcesContext.ResourcesContextProvider>
-        </WorkspaceContextProvider>
+        <resourcesContext.ResourcesContextProvider>
+          <ConnectMyComputerContextProvider rootClusterUri={rootCluster.uri}>
+            {children}
+          </ConnectMyComputerContextProvider>
+        </resourcesContext.ResourcesContextProvider>
       </MockAppContextProvider>
     ),
   });
@@ -322,15 +318,13 @@ describe('canUse', () => {
       const { result } = renderHook(() => useConnectMyComputerContext(), {
         wrapper: ({ children }) => (
           <MockAppContextProvider appContext={appContext}>
-            <WorkspaceContextProvider value={null}>
-              <resourcesContext.ResourcesContextProvider>
-                <ConnectMyComputerContextProvider
-                  rootClusterUri={rootCluster.uri}
-                >
-                  {children}
-                </ConnectMyComputerContextProvider>
-              </resourcesContext.ResourcesContextProvider>
-            </WorkspaceContextProvider>
+            <resourcesContext.ResourcesContextProvider>
+              <ConnectMyComputerContextProvider
+                rootClusterUri={rootCluster.uri}
+              >
+                {children}
+              </ConnectMyComputerContextProvider>
+            </resourcesContext.ResourcesContextProvider>
           </MockAppContextProvider>
         ),
       });

web/packages/teleterm/src/ui/Documents/workspaceContext.tsx

@@ -16,21 +16,27 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import React, { PropsWithChildren } from 'react';
+import {
+  FC,
+  PropsWithChildren,
+  useCallback,
+  useContext,
+  createContext,
+} from 'react';
 
 import { DocumentsService } from 'teleterm/ui/services/workspacesService';
 import { AccessRequestsService } from 'teleterm/ui/services/workspacesService/accessRequestsService';
-import { useAppContext } from 'teleterm/ui/appContextProvider';
 import { ClusterUri, RootClusterUri } from 'teleterm/ui/uri';
+import { useStoreSelector } from 'teleterm/ui/hooks/useStoreSelector';
 
-const WorkspaceContext = React.createContext<{
+const WorkspaceContext = createContext<{
   rootClusterUri: RootClusterUri;
   localClusterUri: ClusterUri;
   documentsService: DocumentsService;
   accessRequestsService: AccessRequestsService;
 }>(null);
 
-export const WorkspaceContextProvider: React.FC<
+export const WorkspaceContextProvider: FC<
   PropsWithChildren<{
     value: {
       rootClusterUri: RootClusterUri;
@@ -40,12 +46,29 @@ export const WorkspaceContextProvider: React.FC<
     };
   }>
 > = props => {
+  // Re-render the context provider whenever the state of the relevant workspace changes. The
+  // context provider cannot re-render only when its props change.
+  // For example, if a new document gets added, none of the props are going to change, but the
+  // callsite that uses useWorkspaceContext might want to get re-rendered in this case, as
+  // technically documentsService returned from useWorkspaceContext might return new state.
+  useStoreSelector(
+    'workspacesService',
+    useCallback(
+      state => state.workspaces[props.value.rootClusterUri],
+      [props.value.rootClusterUri]
+    )
+  );
   return <WorkspaceContext.Provider {...props} />;
 };
 
 export const useWorkspaceContext = () => {
-  const ctx = useAppContext();
-  ctx.workspacesService.useState();
+  const context = useContext(WorkspaceContext);
 
-  return React.useContext(WorkspaceContext);
+  if (!context) {
+    throw new Error(
+      'useWorkspaceContext must be used within a WorkspaceContextProvider'
+    );
+  }
+
+  return context;
 };

web/packages/teleterm/src/ui/Search/SearchBar.tsx

@@ -16,7 +16,7 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import React, { useRef, useEffect } from 'react';
+import React, { useRef, useEffect, useCallback } from 'react';
 import styled from 'styled-components';
 import { Box, Flex } from 'design';
 
@@ -31,14 +31,17 @@ import {
 } from 'teleterm/ui/services/keyboardShortcuts';
 
 import { useAppContext } from '../appContextProvider';
+import { useStoreSelector } from '../hooks/useStoreSelector';
 
 const OPEN_SEARCH_BAR_SHORTCUT_ACTION: KeyboardShortcutAction = 'openSearchBar';
 
 export function SearchBarConnected() {
-  const { workspacesService } = useAppContext();
-  workspacesService.useState();
+  const rootClusterUri = useStoreSelector(
+    'workspacesService',
+    useCallback(state => state.rootClusterUri, [])
+  );
 
-  if (!workspacesService.getRootClusterUri()) {
+  if (!rootClusterUri) {
     return null;
   }
 

web/packages/teleterm/src/ui/Search/SearchContext.tsx

@@ -33,6 +33,7 @@ import { useAppContext } from 'teleterm/ui/appContextProvider';
 import {
   Document,
   DocumentClusterQueryParams,
+  useWorkspaceServiceState,
 } from 'teleterm/ui/services/workspacesService';
 
 import { actionPicker, SearchPicker } from './pickers/pickers';
@@ -130,7 +131,7 @@ export const SearchContextProvider: FC<PropsWithChildren> = props => {
     );
   }
 
-  appContext.workspacesService.useState();
+  useWorkspaceServiceState();
   const activeDocument = appContext.workspacesService
     .getActiveWorkspaceDocumentService()
     ?.getActive();

web/packages/teleterm/src/ui/Search/pickers/useDisplayResults.ts

@@ -25,13 +25,14 @@ import {
   DisplayResults,
 } from 'teleterm/ui/Search/searchResult';
 import { useAppContext } from 'teleterm/ui/appContextProvider';
+import { useWorkspaceServiceState } from 'teleterm/ui/services/workspacesService';
 
 export function useDisplayResults(args: {
   filters: SearchFilter[];
   inputValue: string;
 }): DisplayResults {
   const { workspacesService } = useAppContext();
-  workspacesService.useState();
+  useWorkspaceServiceState();
 
   const localClusterUri =
     workspacesService.getActiveWorkspace()?.localClusterUri;

web/packages/teleterm/src/ui/StatusBar/ShareFeedback/ShareFeedback.test.tsx

@@ -16,14 +16,17 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import React from 'react';
 import { screen } from '@testing-library/react';
 import { fireEvent, render } from 'design/utils/testing';
 
 import { MockAppContextProvider } from 'teleterm/ui/fixtures/MockAppContextProvider';
 import { MockAppContext } from 'teleterm/ui/fixtures/mocks';
 import { IAppContext } from 'teleterm/ui/types';
-import { Cluster } from 'teleterm/services/tshd/types';
+
+import {
+  makeLoggedInUser,
+  makeRootCluster,
+} from 'teleterm/services/tshd/testHelpers';
 
 import { ShareFeedback } from './ShareFeedback';
 
@@ -41,48 +44,42 @@ function renderOpenedShareFeedback(appContext: IAppContext) {
 test('email field is not prefilled with the username if is not an email', () => {
   const appContext = new MockAppContext();
   const clusterUri = '/clusters/localhost';
-  jest
-    .spyOn(appContext.clustersService, 'findCluster')
-    .mockImplementation(() => {
-      return {
-        loggedInUser: { name: 'alice' },
-      } as Cluster;
-    });
-
-  jest
-    .spyOn(appContext.workspacesService, 'getRootClusterUri')
-    .mockReturnValue(clusterUri);
+  appContext.workspacesService.setState(draft => {
+    draft.rootClusterUri = clusterUri;
+  });
+  appContext.clustersService.setState(draft => {
+    draft.clusters.set(
+      clusterUri,
+      makeRootCluster({
+        uri: clusterUri,
+        loggedInUser: makeLoggedInUser({ name: 'alice' }),
+      })
+    );
+  });
 
   renderOpenedShareFeedback(appContext);
 
-  expect(appContext.clustersService.findCluster).toHaveBeenCalledWith(
-    clusterUri
-  );
   expect(screen.getByLabelText('Email Address')).toHaveValue('');
 });
 
 test('email field is prefilled with the username if it looks like an email', () => {
   const appContext = new MockAppContext();
   const clusterUri = '/clusters/production';
-  jest
-    .spyOn(appContext.clustersService, 'findCluster')
-    .mockImplementation(() => {
-      return {
-        loggedInUser: {
-          name: 'bob@prod.com',
-        },
-      } as Cluster;
-    });
-
-  jest
-    .spyOn(appContext.workspacesService, 'getRootClusterUri')
-    .mockReturnValue(clusterUri);
+  appContext.workspacesService.setState(draft => {
+    draft.rootClusterUri = clusterUri;
+  });
+  appContext.clustersService.setState(draft => {
+    draft.clusters.set(
+      clusterUri,
+      makeRootCluster({
+        uri: clusterUri,
+        loggedInUser: makeLoggedInUser({ name: 'bob@prod.com' }),
+      })
+    );
+  });
 
   renderOpenedShareFeedback(appContext);
 
-  expect(appContext.clustersService.findCluster).toHaveBeenCalledWith(
-    clusterUri
-  );
   expect(screen.getByLabelText('Email Address')).toHaveValue('bob@prod.com');
 });
 

web/packages/teleterm/src/ui/StatusBar/ShareFeedback/useShareFeedback.ts

@@ -16,21 +16,25 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import { useState } from 'react';
+import { useCallback, useState } from 'react';
 
 import { makeEmptyAttempt, useAsync } from 'shared/hooks/useAsync';
 
 import { staticConfig } from 'teleterm/staticConfig';
 
 import { useAppContext } from 'teleterm/ui/appContextProvider';
+import { useStoreSelector } from 'teleterm/ui/hooks/useStoreSelector';
 
 import { ShareFeedbackFormValues } from './types';
 
 export const FEEDBACK_TOO_LONG_ERROR = 'FEEDBACK_TOO_LONG_ERROR';
 
 export function useShareFeedback() {
   const ctx = useAppContext();
-  ctx.workspacesService.useState();
+  const rootClusterUri = useStoreSelector(
+    'workspacesService',
+    useCallback(state => state.rootClusterUri, [])
+  );
   ctx.clustersService.useState();
 
   const [isShareFeedbackOpened, setIsShareFeedbackOpened] = useState(false);
@@ -74,9 +78,7 @@ export function useShareFeedback() {
   }
 
   function getEmailFromUserName(): string {
-    const cluster = ctx.clustersService.findCluster(
-      ctx.workspacesService.getRootClusterUri()
-    );
+    const cluster = ctx.clustersService.findCluster(rootClusterUri);
     const userName = cluster?.loggedInUser?.name;
     if (/^\S+@\S+$/.test(userName)) {
       return userName;

web/packages/teleterm/src/ui/StatusBar/useAccessRequestCheckoutButton.ts

@@ -16,11 +16,13 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
+import { useWorkspaceServiceState } from 'teleterm/ui/services/workspacesService';
+
 import { useAppContext } from '../appContextProvider';
 
 export function useAccessRequestsButton() {
   const ctx = useAppContext();
-  ctx.workspacesService.useState();
+  useWorkspaceServiceState();
 
   const workspaceAccessRequest =
     ctx.workspacesService.getActiveWorkspaceAccessRequestsService();

web/packages/teleterm/src/ui/StatusBar/useActiveDocumentClusterBreadcrumbs.ts

@@ -17,12 +17,15 @@
  */
 
 import { useAppContext } from 'teleterm/ui/appContextProvider';
-import { getResourceUri } from 'teleterm/ui/services/workspacesService';
+import {
+  getResourceUri,
+  useWorkspaceServiceState,
+} from 'teleterm/ui/services/workspacesService';
 import { routing } from 'teleterm/ui/uri';
 
 export function useActiveDocumentClusterBreadcrumbs(): string {
   const ctx = useAppContext();
-  ctx.workspacesService.useState();
+  useWorkspaceServiceState();
   ctx.clustersService.useState();
 
   const activeDocument = ctx.workspacesService