fix(app): Fix failedCommand caching issues

app/src/organisms/ErrorRecoveryFlows/hooks/__tests__/useCurrentlyRecoveringFrom.test.ts

@@ -42,9 +42,11 @@ describe('useCurrentlyRecoveringFrom', () => {
           },
         },
       },
+      isFetching: false,
     } as any)
     vi.mocked(useCommandQuery).mockReturnValue({
       data: { data: 'mockCommandDetails' },
+      isFetching: false,
     } as any)
 
     const { result } = renderHook(() =>
@@ -69,8 +71,11 @@ describe('useCurrentlyRecoveringFrom', () => {
       data: {
         links: {},
       },
+      isFetching: false,
+    } as any)
+    vi.mocked(useCommandQuery).mockReturnValue({
+      isFetching: false,
     } as any)
-    vi.mocked(useCommandQuery).mockReturnValue({} as any)
 
     const { result } = renderHook(() =>
       useCurrentlyRecoveringFrom(MOCK_RUN_ID, RUN_STATUS_AWAITING_RECOVERY)
@@ -94,9 +99,11 @@ describe('useCurrentlyRecoveringFrom', () => {
           },
         },
       },
+      isFetching: false,
     } as any)
     vi.mocked(useCommandQuery).mockReturnValue({
       data: { data: 'mockCommandDetails' },
+      isFetching: false,
     } as any)
 
     const { result } = renderHook(() =>
@@ -111,6 +118,91 @@ describe('useCurrentlyRecoveringFrom', () => {
     expect(result.current).toStrictEqual('mockCommandDetails')
   })
 
+  it('returns null if all commands query is still fetching', () => {
+    vi.mocked(useNotifyAllCommandsQuery).mockReturnValue({
+      data: {
+        links: {
+          currentlyRecoveringFrom: {
+            meta: {
+              runId: MOCK_RUN_ID,
+              commandId: MOCK_COMMAND_ID,
+            },
+          },
+        },
+      },
+      isFetching: true,
+    } as any)
+    vi.mocked(useCommandQuery).mockReturnValue({
+      data: { data: 'mockCommandDetails' },
+      isFetching: false,
+    } as any)
+
+    const { result } = renderHook(() =>
+      useCurrentlyRecoveringFrom(MOCK_RUN_ID, RUN_STATUS_AWAITING_RECOVERY)
+    )
+
+    expect(result.current).toStrictEqual(null)
+  })
+
+  it('returns null if command query is still fetching', () => {
+    vi.mocked(useNotifyAllCommandsQuery).mockReturnValue({
+      data: {
+        links: {
+          currentlyRecoveringFrom: {
+            meta: {
+              runId: MOCK_RUN_ID,
+              commandId: MOCK_COMMAND_ID,
+            },
+          },
+        },
+      },
+      isFetching: false,
+    } as any)
+    vi.mocked(useCommandQuery).mockReturnValue({
+      data: { data: 'mockCommandDetails' },
+      isFetching: true,
+    } as any)
+
+    const { result } = renderHook(() =>
+      useCurrentlyRecoveringFrom(MOCK_RUN_ID, RUN_STATUS_AWAITING_RECOVERY)
+    )
+
+    expect(result.current).toStrictEqual(null)
+  })
+
+  it('resets isReadyToShow when run exits recovery mode', () => {
+    const { rerender, result } = renderHook(
+      ({ status }) => useCurrentlyRecoveringFrom(MOCK_RUN_ID, status),
+      { initialProps: { status: RUN_STATUS_AWAITING_RECOVERY } }
+    )
+
+    vi.mocked(useNotifyAllCommandsQuery).mockReturnValue({
+      data: {
+        links: {
+          currentlyRecoveringFrom: {
+            meta: {
+              runId: MOCK_RUN_ID,
+              commandId: MOCK_COMMAND_ID,
+            },
+          },
+        },
+      },
+      isFetching: false,
+    } as any)
+    vi.mocked(useCommandQuery).mockReturnValue({
+      data: { data: 'mockCommandDetails' },
+      isFetching: false,
+    } as any)
+
+    rerender({ status: RUN_STATUS_AWAITING_RECOVERY })
+
+    expect(result.current).toStrictEqual('mockCommandDetails')
+
+    rerender({ status: RUN_STATUS_IDLE } as any)
+
+    expect(result.current).toStrictEqual(null)
+  })
+
   it('calls invalidateQueries when the run enters recovery mode', () => {
     renderHook(() =>
       useCurrentlyRecoveringFrom(MOCK_RUN_ID, RUN_STATUS_AWAITING_RECOVERY)

app/src/organisms/ErrorRecoveryFlows/hooks/useCurrentlyRecoveringFrom.ts

@@ -1,4 +1,4 @@
-import { useEffect } from 'react'
+import { useEffect, useState } from 'react'
 import { useQueryClient } from 'react-query'
 
 import {
@@ -23,13 +23,15 @@ const VALID_RECOVERY_FETCH_STATUSES = [
 ] as Array<RunStatus | null>
 
 // Return the `currentlyRecoveringFrom` command returned by the server, if any.
-// Otherwise, returns null.
+// The command will only be returned after the initial fetches are complete to prevent rendering of stale data.
 export function useCurrentlyRecoveringFrom(
   runId: string,
   runStatus: RunStatus | null
 ): FailedCommand | null {
   const queryClient = useQueryClient()
   const host = useHost()
+  const [isReadyToShow, setIsReadyToShow] = useState(false)
+
   // There can only be a currentlyRecoveringFrom command when the run is in recovery mode.
   // In case we're falling back to polling, only enable queries when that is the case.
   const isRunInRecoveryMode = VALID_RECOVERY_FETCH_STATUSES.includes(runStatus)
@@ -38,10 +40,15 @@ export function useCurrentlyRecoveringFrom(
   useEffect(() => {
     if (isRunInRecoveryMode) {
       void queryClient.invalidateQueries([host, 'runs', runId])
+    } else {
+      setIsReadyToShow(false)
     }
   }, [isRunInRecoveryMode, host, runId])
 
-  const { data: allCommandsQueryData } = useNotifyAllCommandsQuery(
+  const {
+    data: allCommandsQueryData,
+    isFetching: isAllCommandsFetching,
+  } = useNotifyAllCommandsQuery(
     runId,
     { cursor: null, pageLength: 0 }, // pageLength 0 because we only care about the links.
     {
@@ -54,13 +61,36 @@ export function useCurrentlyRecoveringFrom(
 
   // TODO(mm, 2024-05-21): When the server supports fetching the
   // currentlyRecoveringFrom command in one step, do that instead of this chained query.
-  const { data: commandQueryData } = useCommandQuery(
+  const {
+    data: commandQueryData,
+    isFetching: isCommandFetching,
+  } = useCommandQuery(
     currentlyRecoveringFromLink?.meta.runId ?? null,
     currentlyRecoveringFromLink?.meta.commandId ?? null,
     {
       enabled: currentlyRecoveringFromLink != null && isRunInRecoveryMode,
     }
   )
 
-  return isRunInRecoveryMode ? commandQueryData?.data ?? null : null
+  // Only mark as ready to show when waterfall fetches are complete
+  useEffect(() => {
+    if (
+      isRunInRecoveryMode &&
+      !isAllCommandsFetching &&
+      !isCommandFetching &&
+      !isReadyToShow
+    ) {
+      setIsReadyToShow(true)
+    }
+  }, [
+    isRunInRecoveryMode,
+    isAllCommandsFetching,
+    isCommandFetching,
+    isReadyToShow,
+  ])
+
+  const shouldShowCommand =
+    isRunInRecoveryMode && isReadyToShow && commandQueryData?.data
+
+  return shouldShowCommand ? commandQueryData.data : null
 }

app/src/organisms/ErrorRecoveryFlows/hooks/useERUtils.ts

@@ -1,3 +1,5 @@
+import { useMemo } from 'react'
+
 import { useInstrumentsQuery } from '@opentrons/react-api-client'
 
 import { useRouteUpdateActions } from './useRouteUpdateActions'
@@ -13,12 +15,12 @@ import {
 } from '/app/resources/runs'
 import { useRecoveryOptionCopy } from './useRecoveryOptionCopy'
 import { useRecoveryActionMutation } from './useRecoveryActionMutation'
-import { useRunningStepCounts } from '/app/resources/protocols/hooks'
 import { useRecoveryToasts } from './useRecoveryToasts'
 import { useRecoveryAnalytics } from '/app/redux-resources/analytics'
 import { useShowDoorInfo } from './useShowDoorInfo'
 import { useCleanupRecoveryState } from './useCleanupRecoveryState'
 import { useFailedPipetteUtils } from './useFailedPipetteUtils'
+import { getRunningStepCountsFrom } from '/app/resources/protocols'
 
 import type {
   LabwareDefinition2,
@@ -102,7 +104,14 @@ export function useERUtils({
     pageLength: 999,
   })
 
-  const stepCounts = useRunningStepCounts(runId, runCommands)
+  const stepCounts = useMemo(
+    () =>
+      getRunningStepCountsFrom(
+        protocolAnalysis?.commands ?? [],
+        failedCommand?.byRunRecord ?? null
+      ),
+    [protocolAnalysis != null, failedCommand]
+  )
 
   const analytics = useRecoveryAnalytics()
 

app/src/organisms/ErrorRecoveryFlows/hooks/useRetainedFailedCommandBySource.ts

@@ -1,4 +1,4 @@
-import { useState, useEffect } from 'react'
+import { useState, useLayoutEffect } from 'react'
 
 import type { RunTimeCommand } from '@opentrons/shared-data'
 import type { ErrorRecoveryFlowsProps } from '..'
@@ -27,7 +27,7 @@ export function useRetainedFailedCommandBySource(
     setRetainedFailedCommand,
   ] = useState<FailedCommandBySource | null>(null)
 
-  useEffect(() => {
+  useLayoutEffect(() => {
     if (failedCommandByRunRecord !== null) {
       const failedCommandByAnalysis =
         protocolAnalysis?.commands.find(

app/src/organisms/ErrorRecoveryFlows/index.tsx

@@ -1,4 +1,4 @@
-import { useMemo, useState } from 'react'
+import { useMemo, useEffect, useState } from 'react'
 import { useSelector } from 'react-redux'
 
 import {
@@ -7,10 +7,12 @@ import {
   RUN_STATUS_AWAITING_RECOVERY_PAUSED,
   RUN_STATUS_BLOCKED_BY_OPEN_DOOR,
   RUN_STATUS_FAILED,
+  RUN_STATUS_FINISHING,
   RUN_STATUS_IDLE,
   RUN_STATUS_PAUSED,
   RUN_STATUS_RUNNING,
   RUN_STATUS_STOP_REQUESTED,
+  RUN_STATUS_STOPPED,
   RUN_STATUS_SUCCEEDED,
 } from '@opentrons/api-client'
 import {
@@ -41,18 +43,20 @@ const VALID_ER_RUN_STATUSES: RunStatus[] = [
   RUN_STATUS_STOP_REQUESTED,
 ]
 
+// Effectively statuses that are not an "awaiting-recovery" status OR "stop requested."
 const INVALID_ER_RUN_STATUSES: RunStatus[] = [
   RUN_STATUS_RUNNING,
   RUN_STATUS_PAUSED,
   RUN_STATUS_BLOCKED_BY_OPEN_DOOR,
+  RUN_STATUS_FINISHING,
+  RUN_STATUS_STOPPED,
   RUN_STATUS_FAILED,
   RUN_STATUS_SUCCEEDED,
   RUN_STATUS_IDLE,
 ]
 
 export interface UseErrorRecoveryResult {
   isERActive: boolean
-  /* There is no FailedCommand if the run statis is not AWAITING_RECOVERY. */
   failedCommand: FailedCommand | null
 }
 
@@ -61,45 +65,48 @@ export function useErrorRecoveryFlows(
   runStatus: RunStatus | null
 ): UseErrorRecoveryResult {
   const [isERActive, setIsERActive] = useState(false)
-  // If client accesses a valid ER runs status besides AWAITING_RECOVERY but accesses it outside of Error Recovery flows, don't show ER.
   const [hasSeenAwaitingRecovery, setHasSeenAwaitingRecovery] = useState(false)
   const failedCommand = useCurrentlyRecoveringFrom(runId, runStatus)
 
-  if (
-    !hasSeenAwaitingRecovery &&
-    ([
-      RUN_STATUS_AWAITING_RECOVERY,
-      RUN_STATUS_AWAITING_RECOVERY_BLOCKED_BY_OPEN_DOOR,
-      RUN_STATUS_AWAITING_RECOVERY_PAUSED,
-    ] as Array<RunStatus | null>).includes(runStatus)
-  ) {
-    setHasSeenAwaitingRecovery(true)
-  }
-  // Reset recovery mode after the client has exited recovery, otherwise "cancel run" will trigger ER after the first recovery.
-  else if (
-    hasSeenAwaitingRecovery &&
-    runStatus != null &&
-    INVALID_ER_RUN_STATUSES.includes(runStatus)
-  ) {
-    setHasSeenAwaitingRecovery(false)
-  }
-
-  const isValidRunStatus =
-    runStatus != null &&
-    VALID_ER_RUN_STATUSES.includes(runStatus) &&
-    hasSeenAwaitingRecovery
+  // The complexity of this logic exists to persist Error Recovery screens past the server's definition of Error Recovery.
+  // Ex, show a "cancelling run" modal in Error Recovery flows despite the robot no longer being in a recoverable state.
 
-  if (!isERActive && isValidRunStatus && failedCommand != null) {
-    setIsERActive(true)
-  }
-  // Because multiple ER flows may occur per run, disable ER when the status is not "awaiting-recovery" or a
-  // terminating run status in which we want to persist ER flows. Specific recovery commands cause run status to change.
-  // See a specific command's docstring for details.
-  // ER handles a null failedCommand outside the splash screen, so we shouldn't set it false here.
-  else if (isERActive && !isValidRunStatus) {
-    setIsERActive(false)
+  const isValidERStatus = (status: RunStatus | null): boolean => {
+    return (
+      status !== null &&
+      (status === RUN_STATUS_AWAITING_RECOVERY ||
+        (VALID_ER_RUN_STATUSES.includes(status) && hasSeenAwaitingRecovery))
+    )
   }
 
+  // If client accesses a valid ER runs status besides AWAITING_RECOVERY but accesses it outside of Error Recovery flows,
+  // don't show ER.
+  useEffect(() => {
+    if (runStatus != null) {
+      const isAwaitingRecovery =
+        VALID_ER_RUN_STATUSES.includes(runStatus) &&
+        runStatus !== RUN_STATUS_STOP_REQUESTED
+
+      if (isAwaitingRecovery && !hasSeenAwaitingRecovery) {
+        setHasSeenAwaitingRecovery(true)
+      } else if (INVALID_ER_RUN_STATUSES.includes(runStatus)) {
+        setHasSeenAwaitingRecovery(false)
+      }
+    }
+  }, [runStatus, hasSeenAwaitingRecovery])
+
+  // Manage isERActive state, the condition that actually renders Error Recovery.
+  useEffect(() => {
+    const shouldBeActive =
+      isValidERStatus(runStatus) &&
+      // The failedCommand is null when a stop is requested, but we still want to persist Error Recovery in specific circumstances.
+      (failedCommand !== null || runStatus === RUN_STATUS_STOP_REQUESTED)
+
+    if (shouldBeActive !== isERActive) {
+      setIsERActive(shouldBeActive)
+    }
+  }, [runStatus, failedCommand, hasSeenAwaitingRecovery, isERActive])
+
   return {
     isERActive,
     failedCommand,