Hook it up

web/packages/new/photos/components/MLSettings.tsx

@@ -4,6 +4,8 @@ import {
     enableML,
     getIsMLEnabledRemote,
     isMLEnabled,
+    mlStatusSnapshot,
+    mlStatusSubscribe,
     pauseML,
 } from "@/new/photos/services/ml";
 import { EnteDrawer } from "@/new/shared/components/EnteDrawer";
@@ -69,12 +71,10 @@ export const MLSettings: React.FC<MLSettingsProps> = ({
 
     const [status, setStatus] = useState<Status>("loading");
     const [openFaceConsent, setOpenFaceConsent] = useState(false);
-    const [isEnabledLocal, setIsEnabledLocal] = useState(false);
 
     const refreshStatus = async () => {
         if (isMLEnabled() || (await getIsMLEnabledRemote())) {
             setStatus("enabled");
-            setIsEnabledLocal(isMLEnabled());
         } else if (await canEnableML()) {
             setStatus("disabled");
         } else {
@@ -110,7 +110,6 @@ export const MLSettings: React.FC<MLSettingsProps> = ({
             } else {
                 await enableML();
                 setStatus("enabled");
-                setIsEnabledLocal(isMLEnabled());
             }
         } catch (e) {
             log.error("Failed to enable or resume ML", e);
@@ -125,7 +124,6 @@ export const MLSettings: React.FC<MLSettingsProps> = ({
         try {
             await enableML();
             setStatus("enabled");
-            setIsEnabledLocal(isMLEnabled());
             // Close the FaceConsent drawer, come back to ourselves.
             setOpenFaceConsent(false);
         } catch (e) {
@@ -139,7 +137,6 @@ export const MLSettings: React.FC<MLSettingsProps> = ({
     const handleToggleLocal = async () => {
         try {
             isMLEnabled() ? pauseML() : await handleEnableOrResumeML();
-            setIsEnabledLocal(isMLEnabled());
         } catch (e) {
             log.error("Failed to toggle local state of ML", e);
             somethingWentWrong();
@@ -165,7 +162,7 @@ export const MLSettings: React.FC<MLSettingsProps> = ({
         disabled: <EnableML onEnable={handleEnableOrResumeML} />,
         enabled: (
             <ManageML
-                {...{ isEnabledLocal, setDialogBoxAttributesV2 }}
+                {...{ setDialogBoxAttributesV2 }}
                 onToggleLocal={handleToggleLocal}
                 onDisableML={handleDisableML}
             />
@@ -359,8 +356,6 @@ const FaceConsent: React.FC<FaceConsentProps> = ({
 };
 
 interface ManageMLProps {
-    /** `true` if ML is enabled locally (in addition to remote). */
-    isEnabledLocal: boolean;
     /** Called when the user wants to toggle the ML status locally. */
     onToggleLocal: () => void;
     /** Called when the user wants to disable ML. */
@@ -370,12 +365,15 @@ interface ManageMLProps {
 }
 
 const ManageML: React.FC<ManageMLProps> = ({
-    isEnabledLocal,
     onToggleLocal,
     onDisableML,
     setDialogBoxAttributesV2,
 }) => {
-    const status = useSyncExternalStore();
+    const { phase, nSyncedFiles, nTotalFiles } = useSyncExternalStore(
+        mlStatusSubscribe,
+        mlStatusSnapshot,
+    );
+
     const confirmDisableML = () => {
         setDialogBoxAttributesV2({
             title: pt("Disable ML search"),
@@ -458,7 +456,7 @@ const ManageML: React.FC<ManageMLProps> = ({
                     <EnteMenuItem
                         label={pt("On this device")}
                         variant="toggle"
-                        checked={isEnabledLocal}
+                        checked={phase != "paused"}
                         onClick={onToggleLocal}
                     />
                 </MenuItemGroup>
@@ -480,7 +478,7 @@ const ManageML: React.FC<ManageMLProps> = ({
                     >
                         <Typography color="text.muted">Processed</Typography>
                         <Typography textAlign="right">
-                            33,000,000 / 13,000,000
+                            {`${nSyncedFiles} / ${nTotalFiles}`}
                         </Typography>
                     </Stack>
                 </Stack>

web/packages/new/photos/services/ml/index.ts

@@ -42,12 +42,25 @@ let _comlinkWorker: ComlinkWorker<typeof MLWorker> | undefined;
  */
 let _mlStatusListeners: (() => void)[] = [];
 
+/**
+ * Type-wise, we should''ve used undefined to indicate that we don't yet have a
+ * snapshot, but that would make the {@link mlStatusSnapshot} async,
+ * complicating its usage with React's {@link useSyncExternalStore}.
+ *
+ * So instead this value stands in for an `undefined` {@link MLStatus}.
+ */
+const placeholderMLStatus: MLStatus = {
+    phase: "paused",
+    nSyncedFiles: 0,
+    nTotalFiles: 0,
+};
+
 /**
  * Snapshot of {@link MLStatus}.
  *
  * See {@link mlStatusSnapshot}.
  */
-let _mlStatusSnapshot: MLStatus | undefined;
+let _mlStatusSnapshot = placeholderMLStatus;
 
 /** Lazily created, cached, instance of {@link MLWorker}. */
 const worker = async () => {
@@ -104,7 +117,7 @@ export const logoutML = async () => {
     // function (`logoutML`) gets called at a later point in time.
     _isMLEnabled = false;
     _mlStatusListeners = [];
-    _mlStatusSnapshot = undefined;
+    _mlStatusSnapshot = placeholderMLStatus;
     await clearMLDB();
 };
 
@@ -295,39 +308,41 @@ export interface MLStatus {
  */
 export const mlStatusSubscribe = (onChange: () => void): (() => void) => {
     _mlStatusListeners.push(onChange);
+    // Unconditionally update the snapshot.
+    void updateMLStatusSnapshot();
     return () => {
-        _mlStatusListeners = _mlStatusListeners.filter((v) => v != onChange);
+        _mlStatusListeners = _mlStatusListeners.filter((l) => l != onChange);
     };
 };
 
-export const mlStatusSnapshot = (): MLStatus =>
-    (_mlStatusSnapshot ??= getMLStatus());
+export const mlStatusSnapshot = (): MLStatus => _mlStatusSnapshot;
 
-const getMLStatus = (): MLStatus => {
-    return {
-        phase: "paused",
-        nSyncedFiles: 0,
-        nTotalFiles: 0,
-    };
+export const updateMLStatusSnapshot = async () => {
+    const status = await getMLStatus();
+    _mlStatusSnapshot = status;
+    _mlStatusListeners.forEach((l) => l());
 };
 
 /**
- * Return the current state of the face indexing pipeline.
+ * Return the current state of the ML subsystem.
  *
- * Precondition: ML must be enabled.
+ * Precondition: ML must be enabled on remote, though it is fine if it is paused
+ * locally.
  */
-export const faceIndexingStatus = async (): Promise<MLStatus> => {
-    if (!isMLEnabled())
-        throw new Error("Cannot get indexing status when ML is not enabled");
-
+export const getMLStatus = async (): Promise<MLStatus> => {
     const { indexedCount, indexableCount } = await indexableAndIndexedCounts();
-    const isIndexing = await (await worker()).isIndexing();
 
     let phase: MLStatus["phase"];
-    if (indexableCount > 0) {
-        phase = !isIndexing ? "scheduled" : "indexing";
+    if (!isMLEnabled()) {
+        phase = "paused";
     } else {
-        phase = "done";
+        const isIndexing = await (await worker()).isIndexing();
+
+        if (indexableCount > 0) {
+            phase = !isIndexing ? "scheduled" : "indexing";
+        } else {
+            phase = "done";
+        }
     }
 
     return {