Fix screenshot not available with refetch (#1106)

skyvern-frontend/src/routes/tasks/detail/ActionScreenshot.tsx

@@ -1,29 +1,41 @@
 import { getClient } from "@/api/AxiosClient";
-import { ArtifactApiResponse, ArtifactType } from "@/api/types";
+import { ArtifactApiResponse, ArtifactType, Status } from "@/api/types";
 import { ZoomableImage } from "@/components/ZoomableImage";
 import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 import { useQuery } from "@tanstack/react-query";
 import { useParams } from "react-router-dom";
 import { getImageURL } from "./artifactUtils";
 import { ReloadIcon } from "@radix-ui/react-icons";
+import { statusIsNotFinalized } from "../types";
 
 type Props = {
   stepId: string;
   index: number;
+  taskStatus?: Status; // to give a hint that screenshot may not be available if task is not finalized
 };
 
-function ActionScreenshot({ stepId, index }: Props) {
+function ActionScreenshot({ stepId, index, taskStatus }: Props) {
   const { taskId } = useParams();
   const credentialGetter = useCredentialGetter();
 
-  const { data: artifacts, isFetching } = useQuery<Array<ArtifactApiResponse>>({
+  const { data: artifacts, isLoading } = useQuery<Array<ArtifactApiResponse>>({
     queryKey: ["task", taskId, "steps", stepId, "artifacts"],
     queryFn: async () => {
       const client = await getClient(credentialGetter);
       return client
         .get(`/tasks/${taskId}/steps/${stepId}/artifacts`)
         .then((response) => response.data);
     },
+    refetchInterval: (query) => {
+      const data = query.state.data;
+      const screenshot = data?.filter(
+        (artifact) => artifact.artifact_type === ArtifactType.ActionScreenshot,
+      )?.[index];
+      if (!screenshot) {
+        return 5000;
+      }
+      return false;
+    },
   });
 
   const actionScreenshots = artifacts?.filter(
@@ -32,7 +44,7 @@ function ActionScreenshot({ stepId, index }: Props) {
 
   const screenshot = actionScreenshots?.[index];
 
-  if (isFetching) {
+  if (isLoading) {
     return (
       <div className="mx-auto flex max-h-[400px] flex-col items-center gap-2 overflow-hidden">
         <ReloadIcon className="h-6 w-6 animate-spin" />
@@ -41,12 +53,22 @@ function ActionScreenshot({ stepId, index }: Props) {
     );
   }
 
-  return screenshot ? (
+  if (
+    !screenshot &&
+    taskStatus &&
+    statusIsNotFinalized({ status: taskStatus })
+  ) {
+    return <div>The screenshot for this action is not available yet.</div>;
+  }
+
+  if (!screenshot) {
+    return <div>No screenshot found for this action.</div>;
+  }
+
+  return (
     <figure className="mx-auto flex max-w-full flex-col items-center gap-2 overflow-hidden">
       <ZoomableImage src={getImageURL(screenshot)} alt="llm-screenshot" />
     </figure>
-  ) : (
-    <div>Screenshot not found</div>
   );
 }
 

skyvern-frontend/src/routes/tasks/detail/ScrollableActionList.tsx

@@ -10,14 +10,12 @@ import {
   DotFilledIcon,
 } from "@radix-ui/react-icons";
 import { useQueryClient } from "@tanstack/react-query";
-import { ReactNode, useEffect, useRef } from "react";
+import { ReactNode, useRef } from "react";
 import { useParams } from "react-router-dom";
 import { ActionTypePill } from "./ActionTypePill";
 
 type Props = {
   data: Array<Action | null>;
-  onNext: () => void;
-  onPrevious: () => void;
   onActiveIndexChange: (index: number | "stream") => void;
   activeIndex: number | "stream";
   showStreamOption: boolean;
@@ -42,35 +40,19 @@ function ScrollableActionList({
     Array.from({ length: data.length + 1 }),
   );
 
-  useEffect(() => {
-    if (typeof activeIndex === "number" && refs.current[activeIndex]) {
-      refs.current[activeIndex]?.scrollIntoView({
-        behavior: "smooth",
-        block: "nearest",
-      });
-    }
-    if (activeIndex === "stream") {
-      refs.current[data.length]?.scrollIntoView({
-        behavior: "smooth",
-        block: "nearest",
-      });
-    }
-  }, [activeIndex, data.length]);
-
   function getReverseActions() {
     const elements: ReactNode[] = [];
     for (let i = data.length - 1; i >= 0; i--) {
       const action = data[i];
-      const actionIndex = data.length - i - 1;
       if (!action) {
         continue;
       }
-      const selected = activeIndex === actionIndex;
+      const selected = activeIndex === i;
       elements.push(
         <div
           key={i}
           ref={(element) => {
-            refs.current[actionIndex] = element;
+            refs.current[i] = element;
           }}
           className={cn(
             "flex cursor-pointer rounded-lg border-2 bg-slate-elevation3 hover:border-slate-50",
@@ -80,7 +62,7 @@ function ScrollableActionList({
               "border-slate-50": selected,
             },
           )}
-          onClick={() => onActiveIndexChange(actionIndex)}
+          onClick={() => onActiveIndexChange(i)}
           onMouseEnter={() => {
             queryClient.prefetchQuery({
               queryKey: ["task", taskId, "steps", action.stepId, "artifacts"],

skyvern-frontend/src/routes/tasks/detail/TaskActions.tsx

@@ -9,7 +9,11 @@ import { envCredential } from "@/util/env";
 import { keepPreviousData, useQuery } from "@tanstack/react-query";
 import { useEffect, useState } from "react";
 import { useParams } from "react-router-dom";
-import { statusIsNotFinalized, statusIsRunningOrQueued } from "../types";
+import {
+  statusIsFinalized,
+  statusIsNotFinalized,
+  statusIsRunningOrQueued,
+} from "../types";
 import { ActionScreenshot } from "./ActionScreenshot";
 import { useActions } from "./hooks/useActions";
 import { ScrollableActionList } from "./ScrollableActionList";
@@ -33,7 +37,9 @@ function TaskActions() {
   const { taskId } = useParams();
   const credentialGetter = useCredentialGetter();
   const [streamImgSrc, setStreamImgSrc] = useState<string>("");
-  const [selectedAction, setSelectedAction] = useState<number | "stream">(0);
+  const [selectedAction, setSelectedAction] = useState<
+    number | "stream" | null
+  >(null);
   const costCalculator = useCostCalculator();
 
   const { data: task, isLoading: taskIsLoading } = useQuery<TaskApiResponse>({
@@ -89,7 +95,6 @@ function TaskActions() {
             message.status === "terminated"
           ) {
             socket?.close();
-            setSelectedAction(0);
             if (
               message.status === "failed" ||
               message.status === "terminated"
@@ -126,12 +131,6 @@ function TaskActions() {
     };
   }, [credentialGetter, taskId, taskIsRunningOrQueued]);
 
-  useEffect(() => {
-    if (!taskIsLoading && taskIsNotFinalized) {
-      setSelectedAction("stream");
-    }
-  }, [taskIsLoading, taskIsNotFinalized]);
-
   const { data: steps, isLoading: stepsIsLoading } = useQuery<
     Array<StepApiResponse>
   >({
@@ -165,9 +164,23 @@ function TaskActions() {
     );
   }
 
+  function getActiveSelection() {
+    if (selectedAction === null) {
+      if (taskIsNotFinalized) {
+        return "stream";
+      }
+      return actions.length - 1;
+    }
+    if (selectedAction === "stream" && task && statusIsFinalized(task)) {
+      return actions.length - 1;
+    }
+    return selectedAction;
+  }
+
+  const activeSelection = getActiveSelection();
+
   const activeAction =
-    typeof selectedAction === "number" &&
-    actions?.[actions.length - selectedAction - 1];
+    activeSelection !== "stream" ? actions[activeSelection] : null;
 
   function getStream() {
     if (task?.status === Status.Created) {
@@ -212,17 +225,18 @@ function TaskActions() {
     <div className="flex gap-2">
       <div className="w-2/3 rounded border">
         <div className="h-full w-full p-4">
-          {selectedAction === "stream" ? getStream() : null}
-          {typeof selectedAction === "number" && activeAction ? (
+          {activeSelection === "stream" ? getStream() : null}
+          {typeof activeSelection === "number" && activeAction ? (
             <ActionScreenshot
               stepId={activeAction.stepId}
               index={activeAction.index}
+              taskStatus={task?.status}
             />
           ) : null}
         </div>
       </div>
       <ScrollableActionList
-        activeIndex={selectedAction}
+        activeIndex={activeSelection}
         data={actions ?? []}
         onActiveIndexChange={setSelectedAction}
         showStreamOption={Boolean(taskIsNotFinalized)}
@@ -233,52 +247,6 @@ function TaskActions() {
             ? formatter.format(costCalculator(notRunningSteps ?? []))
             : undefined,
         }}
-        onNext={() => {
-          if (!actions) {
-            return;
-          }
-          setSelectedAction((prev) => {
-            if (taskIsNotFinalized) {
-              if (actions.length === 0) {
-                return "stream";
-              }
-              if (prev === actions.length - 1) {
-                return actions.length - 1;
-              }
-              if (prev === "stream") {
-                return 0;
-              }
-              return prev + 1;
-            }
-            if (typeof prev === "number") {
-              return prev === actions.length - 1 ? prev : prev + 1;
-            }
-            return 0;
-          });
-        }}
-        onPrevious={() => {
-          if (!actions) {
-            return;
-          }
-          setSelectedAction((prev) => {
-            if (taskIsNotFinalized) {
-              if (actions.length === 0) {
-                return "stream";
-              }
-              if (prev === 0) {
-                return "stream";
-              }
-              if (prev === "stream") {
-                return "stream";
-              }
-              return prev - 1;
-            }
-            if (typeof prev === "number") {
-              return prev === 0 ? prev : prev - 1;
-            }
-            return 0;
-          });
-        }}
       />
     </div>
   );

skyvern-frontend/src/routes/tasks/types.ts

@@ -24,6 +24,16 @@ export function statusIsNotFinalized({ status }: { status: Status }): boolean {
   );
 }
 
+export function statusIsFinalized({ status }: { status: Status }): boolean {
+  return (
+    status === Status.Completed ||
+    status === Status.Failed ||
+    status === Status.Terminated ||
+    status === Status.TimedOut ||
+    status === Status.Canceled
+  );
+}
+
 export function statusIsRunningOrQueued({
   status,
 }: {