add useActivatePartnerTheme hook

src/background/initTheme.ts

@@ -42,6 +42,7 @@ async function setToolbarIcon(): Promise<void> {
   }
 
   const activeTheme = await getActiveTheme();
+
   await allSettled(
     [themeStorage.set(activeTheme), setToolbarIconFromTheme(activeTheme)],
     { catch: "ignore" },

src/extensionConsole/pages/onboarding/partner/PartnerSetupCard.tsx

@@ -23,13 +23,13 @@ import ControlRoomOAuthForm from "@/extensionConsole/pages/onboarding/partner/Co
 import ControlRoomTokenForm from "@/extensionConsole/pages/onboarding/partner/ControlRoomTokenForm";
 import { selectSettings } from "@/store/settings/settingsSelectors";
 import { useGetMeQuery } from "@/data/service/api";
-import { useDispatch, useSelector } from "react-redux";
+import { useSelector } from "react-redux";
 import { selectIsLoggedIn } from "@/auth/authSelectors";
 import { Button } from "react-bootstrap";
 import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import { faLink } from "@fortawesome/free-solid-svg-icons";
 import { getBaseURL } from "@/data/service/baseService";
-import { updateLocalPartnerTheme } from "@/store/settings/settingsSlice";
+import { useActivatePartnerTheme } from "@/store/settings/settingsSlice";
 import { useLocation } from "react-router";
 import {
   hostnameToUrl,
@@ -40,7 +40,6 @@ import useManagedStorageState from "@/store/enterprise/useManagedStorageState";
 import { type FetchableAsyncState } from "@/types/sliceTypes";
 import useLinkState from "@/auth/useLinkState";
 import Loader from "@/components/Loader";
-import { activateTheme } from "@/background/messenger/api";
 
 /**
  * Create the app URL for the partner start page. It shows content based on whether or not the hostname corresponds
@@ -114,7 +113,7 @@ function usePartnerLoginMode(): "token" | "oauth2" {
  * Currently, supports the Automation Anywhere partner integration.
  */
 const PartnerSetupCard: React.FunctionComponent = () => {
-  const dispatch = useDispatch();
+  const activatePartnerTheme = useActivatePartnerTheme();
   // Make sure to use useLocation because the location.search are on the hash route
   const location = useLocation();
   const mode = usePartnerLoginMode();
@@ -150,10 +149,8 @@ const PartnerSetupCard: React.FunctionComponent = () => {
   };
 
   useEffect(() => {
-    // Ensure the partner branding is applied
-    dispatch(updateLocalPartnerTheme("automation-anywhere"));
-    void activateTheme();
-  }, [dispatch]);
+    activatePartnerTheme("automation-anywhere");
+  }, [activatePartnerTheme]);
 
   if (isLinkedLoading || isMeLoading) {
     return <Loader />;

src/extensionConsole/pages/settings/AdvancedSettings.tsx

@@ -25,7 +25,7 @@ import { clearCachedAuthSecrets, clearPartnerAuth } from "@/auth/authStorage";
 import notify from "@/utils/notify";
 import useFlags from "@/hooks/useFlags";
 import settingsSlice, {
-  updateLocalPartnerTheme,
+  useActivatePartnerTheme,
 } from "@/store/settings/settingsSlice";
 import { useDispatch, useSelector } from "react-redux";
 import { assertProtocolUrl } from "@/utils/urlUtils";
@@ -51,6 +51,7 @@ const AdvancedSettings: React.FunctionComponent = () => {
   const { partnerId, authIntegrationId, authMethod } =
     useSelector(selectSettings);
   const { exportDiagnostics } = useDiagnostics();
+  const activatePartnerTheme = useActivatePartnerTheme();
 
   const [serviceURL, setServiceURL] = useConfiguredHost();
 
@@ -200,7 +201,7 @@ const AdvancedSettings: React.FunctionComponent = () => {
               placeholder="my-company"
               defaultValue={partnerId ?? ""}
               onBlur={(event: React.FocusEvent<HTMLInputElement>) => {
-                dispatch(updateLocalPartnerTheme(event.target.value));
+                activatePartnerTheme(event.target.value);
               }}
             />
             <Form.Text muted>The partner id of a PixieBrix partner</Form.Text>

src/store/settings/settingsSlice.ts

@@ -30,6 +30,9 @@ import { type RegistryId } from "@/types/registryTypes";
 import { isRegistryId } from "@/types/helpers";
 import { revertAll } from "@/store/commonActions";
 import { activateTheme } from "@/background/messenger/api";
+import { useDispatch, useSelector } from "react-redux";
+import { useCallback, useEffect } from "react";
+import { selectSettings } from "@/store/settings/settingsSelectors";
 
 export const initialSettingsState: SettingsState = {
   nextUpdate: null,
@@ -140,4 +143,24 @@ export const updateLocalPartnerTheme = createAsyncThunk<
   await activateTheme();
 });
 
+export const useActivatePartnerTheme = (): ((
+  partnerId: string | null,
+) => void) => {
+  const dispatch = useDispatch();
+  const { partnerId } = useSelector(selectSettings);
+
+  useEffect(() => {
+    if (partnerId) {
+      void activateTheme();
+    }
+  }, [partnerId]);
+
+  return useCallback(
+    (partnerId: string) => {
+      dispatch(settingsSlice.actions.setPartnerId({ partnerId }));
+    },
+    [dispatch],
+  );
+};
+
 export default settingsSlice;