Fix screenshot not available with refetch (#1106)

2024-11-01 09:00:41 -07:00
parent ac2905f790
commit 8fa1be29e6
4 changed files with 71 additions and 89 deletions
--- a/skyvern-frontend/src/routes/tasks/detail/ActionScreenshot.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/ActionScreenshot.tsx
@@ -1,22 +1,24 @@
 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);
@@ -24,6 +26,16 @@ function ActionScreenshot({ stepId, index }: Props) {
        .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>
  );
 }

--- a/skyvern-frontend/src/routes/tasks/detail/ScrollableActionList.tsx
+++ b/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"],
--- a/skyvern-frontend/src/routes/tasks/detail/TaskActions.tsx
+++ b/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>
  );
--- a/skyvern-frontend/src/routes/tasks/types.ts
+++ b/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,
 }: {