feat(app,robot-server): Subscribe to currentlyRecoveringFrom notifications

app/src/organisms/ErrorRecoveryFlows/__tests__/ErrorRecoveryFlows.test.tsx

@@ -12,7 +12,7 @@ import { renderWithProviders } from '../../../__testing-utils__'
 import { i18n } from '../../../i18n'
 import { mockFailedCommand } from '../__fixtures__'
 import { ErrorRecoveryFlows, useErrorRecoveryFlows } from '..'
-import { useCurrentlyFailedRunCommand } from '../utils'
+import { useCurrentlyRecoveringFrom } from '../utils'
 import { useFeatureFlag } from '../../../redux/config'
 import { useERWizard, ErrorRecoveryWizard } from '../ErrorRecoveryWizard'
 import { useRunPausedSplash, RunPausedSplash } from '../RunPausedSplash'
@@ -27,9 +27,7 @@ vi.mock('../RunPausedSplash')
 
 describe('useErrorRecoveryFlows', () => {
   beforeEach(() => {
-    vi.mocked(useCurrentlyFailedRunCommand).mockReturnValue(
-      'mockCommand' as any
-    )
+    vi.mocked(useCurrentlyRecoveringFrom).mockReturnValue('mockCommand' as any)
   })
 
   it('should have initial state of isEREnabled as false', () => {

app/src/organisms/ErrorRecoveryFlows/__tests__/utils.test.ts

@@ -1,23 +1,24 @@
 import { vi, describe, it, expect, beforeEach } from 'vitest'
-import { renderHook, act } from '@testing-library/react'
-
-import {
-  RUN_STATUS_AWAITING_RECOVERY,
-  RUN_STATUS_RUNNING,
-} from '@opentrons/api-client'
+import { renderHook } from '@testing-library/react'
 
 import { ERROR_KINDS, INVALID, RECOVERY_MAP } from '../constants'
 import {
   getErrorKind,
   getRecoveryRouteNavigation,
   useRouteUpdateActions,
-  useCurrentlyFailedRunCommand,
+  useCurrentlyRecoveringFrom,
 } from '../utils'
 import { useNotifyAllCommandsQuery } from '../../../resources/runs'
 
 import type { Mock } from 'vitest'
 import type { GetRouteUpdateActionsParams } from '../utils'
+import { useCommandQuery } from '@opentrons/react-api-client'
+import {
+  RUN_STATUS_AWAITING_RECOVERY,
+  RUN_STATUS_IDLE,
+} from '@opentrons/api-client'
 
+vi.mock('@opentrons/react-api-client')
 vi.mock('../../../resources/runs')
 
 describe('getErrorKind', () => {
@@ -181,87 +182,88 @@ describe('useRouteUpdateActions', () => {
   })
 })
 
-const MOCK_COMMANDS_QUERY = {
-  data: {
-    data: [
-      { status: 'failed', intent: 'fixit', id: '0' },
-      { status: 'failed', intent: 'protocol', id: '111' },
-      { status: 'failed', intent: 'protocol', id: '123' },
-      { status: 'success', intent: 'fixit', id: '1' },
-    ],
-  },
-} as any
-
 const MOCK_RUN_ID = 'runId'
+const MOCK_COMMAND_ID = 'commandId'
 
-describe('useCurrentlyFailedRunCommand', () => {
-  beforeEach(() => {
-    vi.mocked(useNotifyAllCommandsQuery).mockReturnValue(MOCK_COMMANDS_QUERY)
-  })
+describe('useCurrentlyRecoveringFrom', () => {
+  it('disables all queries if the run is not awaiting-recovery', () => {
+    vi.mocked(useNotifyAllCommandsQuery).mockReturnValue({
+      data: {
+        links: {
+          currentlyRecoveringFrom: {
+            meta: {
+              runId: MOCK_RUN_ID,
+              commandId: MOCK_COMMAND_ID,
+            },
+          },
+        },
+      },
+    } as any)
+    vi.mocked(useCommandQuery).mockReturnValue({
+      data: { data: 'mockCommandDetails' },
+    } as any)
 
-  it('returns null on initial render when the run status is not "awaiting-recovery"', () => {
     const { result } = renderHook(() =>
-      useCurrentlyFailedRunCommand(MOCK_RUN_ID, RUN_STATUS_RUNNING)
+      useCurrentlyRecoveringFrom(MOCK_RUN_ID, RUN_STATUS_IDLE)
     )
 
-    expect(result.current).toBeNull()
+    expect(vi.mocked(useNotifyAllCommandsQuery)).toHaveBeenCalledWith(
+      MOCK_RUN_ID,
+      { cursor: null, pageLength: 0 },
+      { enabled: false, refetchInterval: 5000 }
+    )
+    expect(vi.mocked(useCommandQuery)).toHaveBeenCalledWith(
+      MOCK_RUN_ID,
+      MOCK_COMMAND_ID,
+      { enabled: false }
+    )
+    expect(result.current).toStrictEqual(null)
   })
 
-  it('sets recentFailedCommand correctly when runStatus is "awaiting-recovery" and there is no recent failed command', () => {
-    const { result, rerender } = renderHook(
-      // @ts-expect-error this works
-      props => useCurrentlyFailedRunCommand(...props),
-      {
-        initialProps: [MOCK_RUN_ID, RUN_STATUS_RUNNING],
-      }
-    )
+  it('returns null if there is no currentlyRecoveringFrom command', () => {
+    vi.mocked(useNotifyAllCommandsQuery).mockReturnValue({
+      data: {
+        links: {},
+      },
+    } as any)
+    vi.mocked(useCommandQuery).mockReturnValue({} as any)
 
-    act(() => {
-      rerender([MOCK_RUN_ID, RUN_STATUS_AWAITING_RECOVERY])
-    })
+    const { result } = renderHook(() =>
+      useCurrentlyRecoveringFrom(MOCK_RUN_ID, RUN_STATUS_AWAITING_RECOVERY)
+    )
 
-    expect(result.current).toEqual({
-      status: 'failed',
-      intent: 'protocol',
-      id: '123',
+    expect(vi.mocked(useCommandQuery)).toHaveBeenCalledWith(null, null, {
+      enabled: false,
     })
+    expect(result.current).toStrictEqual(null)
   })
 
-  it('always returns the  failed protocol run command that caused the run to enter "awaiting-recovery"', () => {
-    const { result, rerender } = renderHook(
-      // @ts-expect-error this works
-      props => useCurrentlyFailedRunCommand(...props),
-      {
-        initialProps: [MOCK_RUN_ID, RUN_STATUS_AWAITING_RECOVERY],
-      }
-    )
-
+  it('fetches and returns the currentlyRecoveringFrom command, given that there is one', () => {
     vi.mocked(useNotifyAllCommandsQuery).mockReturnValue({
-      ...MOCK_COMMANDS_QUERY,
-      ...{ status: 'failed', intent: 'protocol', id: '124' },
-    })
-    rerender([MOCK_RUN_ID, RUN_STATUS_AWAITING_RECOVERY])
-
-    expect(result.current).toEqual({
-      status: 'failed',
-      intent: 'protocol',
-      id: '123',
-    })
-  })
+      data: {
+        links: {
+          currentlyRecoveringFrom: {
+            meta: {
+              runId: MOCK_RUN_ID,
+              commandId: MOCK_COMMAND_ID,
+            },
+          },
+        },
+      },
+    } as any)
+    vi.mocked(useCommandQuery).mockReturnValue({
+      data: { data: 'mockCommandDetails' },
+    } as any)
 
-  it('sets recentFailedCommand to null when runStatus is not "awaiting-recovery"', () => {
-    const { result, rerender } = renderHook(
-      // @ts-expect-error this works
-      props => useCurrentlyFailedRunCommand(...props),
-      {
-        initialProps: ['runId', 'awaiting-recovery'],
-      }
+    const { result } = renderHook(() =>
+      useCurrentlyRecoveringFrom(MOCK_RUN_ID, RUN_STATUS_AWAITING_RECOVERY)
     )
 
-    act(() => {
-      rerender([MOCK_RUN_ID, RUN_STATUS_RUNNING])
-    })
-
-    expect(result.current).toBeNull()
+    expect(vi.mocked(useCommandQuery)).toHaveBeenCalledWith(
+      MOCK_RUN_ID,
+      MOCK_COMMAND_ID,
+      { enabled: true }
+    )
+    expect(result.current).toStrictEqual('mockCommandDetails')
   })
 })

app/src/organisms/ErrorRecoveryFlows/index.tsx

@@ -8,7 +8,7 @@ import {
 import { useFeatureFlag } from '../../redux/config'
 import { ErrorRecoveryWizard, useERWizard } from './ErrorRecoveryWizard'
 import { useRunPausedSplash, RunPausedSplash } from './RunPausedSplash'
-import { useCurrentlyFailedRunCommand, useRouteUpdateActions } from './utils'
+import { useCurrentlyRecoveringFrom, useRouteUpdateActions } from './utils'
 import { useRecoveryCommands } from './useRecoveryCommands'
 import { RECOVERY_MAP } from './constants'
 
@@ -30,7 +30,7 @@ export function useErrorRecoveryFlows(
   runStatus: RunStatus | null
 ): UseErrorRecoveryResult {
   const [isERActive, setIsERActive] = React.useState(false)
-  const failedCommand = useCurrentlyFailedRunCommand(runId, runStatus)
+  const failedCommand = useCurrentlyRecoveringFrom(runId, runStatus)
 
   const isValidRunStatus =
     runStatus != null && VALID_ER_RUN_STATUSES.includes(runStatus)

app/src/organisms/ErrorRecoveryFlows/utils.ts

@@ -2,9 +2,9 @@ import * as React from 'react'
 import { useTranslation } from 'react-i18next'
 import head from 'lodash/head'
 import last from 'lodash/last'
-import findLast from 'lodash/findLast'
 
 import { RUN_STATUS_AWAITING_RECOVERY } from '@opentrons/api-client'
+import { useCommandQuery } from '@opentrons/react-api-client'
 
 import { useNotifyAllCommandsQuery } from '../../resources/runs'
 import { RECOVERY_MAP, ERROR_KINDS, INVALID, STEP_ORDER } from './constants'
@@ -21,46 +21,41 @@ import type {
 
 // TODO(jh, 05-09-24): Migrate utils, useRecoveryCommands.ts, and respective tests to a utils dir, and make each util a separate file.
 
-// While the run is "awaiting-recovery", return the most recently failed run command with a protocol intent.
-// Otherwise, returns null.
 const ALL_COMMANDS_POLL_MS = 5000
 
-// TODO(jh, 05-20-24): Update the logic for returning the failed run command once EXEC-458 merges.
-export function useCurrentlyFailedRunCommand(
+// Return the `currentlyRecoveringFrom` command returned by the server, if any.
+// Otherwise, returns null.
+export function useCurrentlyRecoveringFrom(
   runId: string,
   runStatus: RunStatus | null
 ): FailedCommand | null {
-  const [
-    recentFailedCommand,
-    setRecentFailedCommand,
-  ] = React.useState<FailedCommand | null>(null)
-  // The most recently failed protocol command causes the run to enter "awaiting-recovery", therefore only check
-  // for a newly failed command when the run first enters "awaiting-recovery."
+  // There can only be a currentlyRecoveringFrom command when the run is awaiting-recovery.
+  // In case we're falling back to polling, only enable queries when that is the case.
   const isRunStatusAwaitingRecovery = runStatus === RUN_STATUS_AWAITING_RECOVERY
 
   const { data: allCommandsQueryData } = useNotifyAllCommandsQuery(
     runId,
-    null,
+    { cursor: null, pageLength: 0 }, // pageLength 0 because we only care about the links.
     {
-      enabled: isRunStatusAwaitingRecovery && recentFailedCommand == null,
+      enabled: isRunStatusAwaitingRecovery,
       refetchInterval: ALL_COMMANDS_POLL_MS,
     }
   )
-
-  React.useEffect(() => {
-    if (isRunStatusAwaitingRecovery && recentFailedCommand == null) {
-      const failedCommand =
-        findLast(
-          allCommandsQueryData?.data,
-          command => command.status === 'failed' && command.intent !== 'fixit'
-        ) ?? null
-      setRecentFailedCommand(failedCommand)
-    } else if (!isRunStatusAwaitingRecovery && recentFailedCommand != null) {
-      setRecentFailedCommand(null)
+  const currentlyRecoveringFromLink =
+    allCommandsQueryData?.links.currentlyRecoveringFrom
+
+  // 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(
+    currentlyRecoveringFromLink?.meta.runId ?? null,
+    currentlyRecoveringFromLink?.meta.commandId ?? null,
+    {
+      enabled:
+        currentlyRecoveringFromLink != null && isRunStatusAwaitingRecovery,
     }
-  }, [isRunStatusAwaitingRecovery, recentFailedCommand, allCommandsQueryData])
+  )
 
-  return recentFailedCommand
+  return isRunStatusAwaitingRecovery ? commandQueryData?.data ?? null : null
 }
 
 export function useErrorName(errorKind: ErrorKind): string {

app/src/redux/shell/types.ts

@@ -138,7 +138,7 @@ export interface RobotMassStorageDeviceRemoved {
 export type NotifyTopic =
   | 'ALL_TOPICS'
   | 'robot-server/maintenance_runs/current_run'
-  | 'robot-server/runs/current_command'
+  | 'robot-server/runs/commands_links'
   | 'robot-server/runs'
   | `robot-server/runs/${string}`
   | 'robot-server/deck_configuration'

app/src/resources/runs/useNotifyAllCommandsQuery.ts

@@ -11,8 +11,14 @@ export function useNotifyAllCommandsQuery<TError = Error>(
   params?: GetCommandsParams | null,
   options: QueryOptionsWithPolling<CommandsData, TError> = {}
 ): UseQueryResult<CommandsData, TError> {
+  // Assume the useAllCommandsQuery() response can only change when the command links change.
+  //
+  // TODO(mm, 2024-05-21): I don't think this is correct. If a command goes from
+  // running to succeeded, that may change the useAllCommandsQuery response, but it
+  // will not necessarily change the command links. We might need an MQTT topic
+  // covering "any change in `GET /runs/{id}/commands`".
   const { notifyOnSettled, shouldRefetch } = useNotifyDataReady({
-    topic: 'robot-server/runs/current_command', // only updates to the current command cause all commands to change
+    topic: 'robot-server/runs/commands_links',
     options,
   })
 

robot-server/robot_server/runs/run_data_manager.py

@@ -194,8 +194,9 @@ async def create(
             created_at=created_at,
             protocol_id=protocol.protocol_id if protocol is not None else None,
         )
-        await self._runs_publisher.initialize(
+        await self._runs_publisher.start_publishing_for_run(
             get_current_command=self.get_current_command,
+            get_recovery_target_command=self.get_recovery_target_command,
             get_state_summary=self._get_good_state_summary,
             run_id=run_id,
         )