From 5ce37dfaaf70410e1fbe44990da74155d1d92c9f Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Tue, 7 May 2024 11:31:05 -0700
Subject: [PATCH] FE Auth vendor changes (#270)

---
 skyvern-frontend/src/api/AxiosClient.ts       | 25 ++++++++++++++++---
 .../src/components/useThemeAsDarkOrLight.ts   | 15 +++++++++++
 .../src/hooks/useCredentialGetter.ts          |  9 +++++++
 skyvern-frontend/src/routes/root/Profile.tsx  | 11 ++------
 .../src/routes/root/RootLayout.tsx            |  8 ++----
 skyvern-frontend/src/routes/root/SideNav.tsx  |  8 +++---
 .../routes/tasks/create/CreateNewTaskForm.tsx |  7 ++++--
 .../src/routes/tasks/detail/StepArtifacts.tsx |  5 +++-
 .../tasks/detail/StepArtifactsLayout.tsx      |  5 +++-
 .../routes/tasks/detail/StepNavigation.tsx    |  5 +++-
 .../src/routes/tasks/detail/TaskDetails.tsx   |  5 +++-
 .../src/routes/tasks/list/TaskList.tsx        |  5 +++-
 .../routes/tasks/running/LatestScreenshot.tsx |  6 ++++-
 .../src/routes/tasks/running/QueuedTasks.tsx  |  6 ++++-
 .../src/routes/tasks/running/RunningTasks.tsx |  5 +++-
 .../src/store/CredentialGetterContext.ts      |  6 +++++
 16 files changed, 99 insertions(+), 32 deletions(-)
 create mode 100644 skyvern-frontend/src/components/useThemeAsDarkOrLight.ts
 create mode 100644 skyvern-frontend/src/hooks/useCredentialGetter.ts
 create mode 100644 skyvern-frontend/src/store/CredentialGetterContext.ts

diff --git a/skyvern-frontend/src/api/AxiosClient.ts b/skyvern-frontend/src/api/AxiosClient.ts
index ae54a783..2d678870 100644
--- a/skyvern-frontend/src/api/AxiosClient.ts
+++ b/skyvern-frontend/src/api/AxiosClient.ts
@@ -18,7 +18,9 @@ export function setAuthorizationHeader(token: string) {
 }
 
 export function removeAuthorizationHeader() {
-  delete client.defaults.headers.common["Authorization"];
+  if (client.defaults.headers.common["Authorization"]) {
+    delete client.defaults.headers.common["Authorization"];
+  }
 }
 
 export function setApiKeyHeader(apiKey: string) {
@@ -26,7 +28,24 @@ export function setApiKeyHeader(apiKey: string) {
 }
 
 export function removeApiKeyHeader() {
-  delete client.defaults.headers.common["X-API-Key"];
+  if (client.defaults.headers.common["X-API-Key"]) {
+    delete client.defaults.headers.common["X-API-Key"];
+  }
 }
 
-export { client, artifactApiClient };
+async function getClient(credentialGetter: CredentialGetter | null) {
+  if (credentialGetter) {
+    removeApiKeyHeader();
+    const credential = await credentialGetter();
+    if (!credential) {
+      console.warn("No credential found");
+      return client;
+    }
+    setAuthorizationHeader(credential);
+  }
+  return client;
+}
+
+export type CredentialGetter = () => Promise<string | null>;
+
+export { getClient, artifactApiClient };
diff --git a/skyvern-frontend/src/components/useThemeAsDarkOrLight.ts b/skyvern-frontend/src/components/useThemeAsDarkOrLight.ts
new file mode 100644
index 00000000..1a8eac93
--- /dev/null
+++ b/skyvern-frontend/src/components/useThemeAsDarkOrLight.ts
@@ -0,0 +1,15 @@
+import { useTheme } from "./useTheme";
+
+function useThemeAsDarkOrLight(): "light" | "dark" {
+  const { theme: baseTheme } = useTheme();
+
+  if (baseTheme === "dark" || baseTheme === "light") {
+    return baseTheme;
+  }
+
+  return window.matchMedia("(prefers-color-scheme: dark)").matches
+    ? "dark"
+    : "light";
+}
+
+export { useThemeAsDarkOrLight };
diff --git a/skyvern-frontend/src/hooks/useCredentialGetter.ts b/skyvern-frontend/src/hooks/useCredentialGetter.ts
new file mode 100644
index 00000000..0b0689bf
--- /dev/null
+++ b/skyvern-frontend/src/hooks/useCredentialGetter.ts
@@ -0,0 +1,9 @@
+import { CredentialGetterContext } from "@/store/CredentialGetterContext";
+import { useContext } from "react";
+
+function useCredentialGetter() {
+  const credentialGetter = useContext(CredentialGetterContext);
+  return credentialGetter;
+}
+
+export { useCredentialGetter };
diff --git a/skyvern-frontend/src/routes/root/Profile.tsx b/skyvern-frontend/src/routes/root/Profile.tsx
index a4e38d9a..00de292c 100644
--- a/skyvern-frontend/src/routes/root/Profile.tsx
+++ b/skyvern-frontend/src/routes/root/Profile.tsx
@@ -3,10 +3,9 @@ import { ExitIcon, PersonIcon } from "@radix-ui/react-icons";
 
 type Props = {
   name: string;
-  onLogout?: () => void;
 };
 
-function Profile({ name, onLogout }: Props) {
+function Profile({ name }: Props) {
   return (
     <div className="flex items-center border-2 p-2 rounded-lg">
       <div className="flex gap-2 items-center">
@@ -14,13 +13,7 @@ function Profile({ name, onLogout }: Props) {
         <p className="w-40 overflow-hidden text-ellipsis">{name}</p>
       </div>
       <div>
-        <Button
-          variant="outline"
-          size="icon"
-          onClick={() => {
-            onLogout?.();
-          }}
-        >
+        <Button variant="outline" size="icon">
           <ExitIcon className="h-4 w-4" />
         </Button>
       </div>
diff --git a/skyvern-frontend/src/routes/root/RootLayout.tsx b/skyvern-frontend/src/routes/root/RootLayout.tsx
index db2d0975..c334cc3f 100644
--- a/skyvern-frontend/src/routes/root/RootLayout.tsx
+++ b/skyvern-frontend/src/routes/root/RootLayout.tsx
@@ -7,11 +7,7 @@ import { Profile } from "./Profile";
 import { useContext } from "react";
 import { UserContext } from "@/store/UserContext";
 
-type Props = {
-  onLogout?: () => void;
-};
-
-function RootLayout({ onLogout }: Props) {
+function RootLayout() {
   const user = useContext(UserContext);
 
   return (
@@ -30,7 +26,7 @@ function RootLayout({ onLogout }: Props) {
           <SideNav />
           {user ? (
             <div className="absolute bottom-2 left-0 w-72 px-6 shrink-0">
-              <Profile name={user.name} onLogout={onLogout} />
+              <Profile name={user.name} />
             </div>
           ) : null}
         </aside>
diff --git a/skyvern-frontend/src/routes/root/SideNav.tsx b/skyvern-frontend/src/routes/root/SideNav.tsx
index 0fe46c5a..c997201e 100644
--- a/skyvern-frontend/src/routes/root/SideNav.tsx
+++ b/skyvern-frontend/src/routes/root/SideNav.tsx
@@ -8,12 +8,12 @@ import { NavLink } from "react-router-dom";
 
 function SideNav() {
   return (
-    <nav className="flex flex-col gap-4">
+    <nav>
       <NavLink
         to="create"
         className={({ isActive }) => {
           return cn(
-            "flex items-center px-6 py-2 hover:bg-primary-foreground rounded-2xl",
+            "flex items-center px-6 py-4 hover:bg-primary-foreground rounded-2xl",
             {
               "bg-primary-foreground": isActive,
             },
@@ -27,7 +27,7 @@ function SideNav() {
         to="tasks"
         className={({ isActive }) => {
           return cn(
-            "flex items-center px-6 py-2 hover:bg-primary-foreground rounded-2xl",
+            "flex items-center px-6 py-4 hover:bg-primary-foreground rounded-2xl",
             {
               "bg-primary-foreground": isActive,
             },
@@ -41,7 +41,7 @@ function SideNav() {
         to="settings"
         className={({ isActive }) => {
           return cn(
-            "flex items-center px-6 py-2 hover:bg-primary-foreground rounded-2xl",
+            "flex items-center px-6 py-4 hover:bg-primary-foreground rounded-2xl",
             {
               "bg-primary-foreground": isActive,
             },
diff --git a/skyvern-frontend/src/routes/tasks/create/CreateNewTaskForm.tsx b/skyvern-frontend/src/routes/tasks/create/CreateNewTaskForm.tsx
index 9776fad3..1b083369 100644
--- a/skyvern-frontend/src/routes/tasks/create/CreateNewTaskForm.tsx
+++ b/skyvern-frontend/src/routes/tasks/create/CreateNewTaskForm.tsx
@@ -21,7 +21,7 @@ import {
 } from "../data/descriptionHelperContent";
 import { Textarea } from "@/components/ui/textarea";
 import { useMutation, useQueryClient } from "@tanstack/react-query";
-import { client } from "@/api/AxiosClient";
+import { getClient } from "@/api/AxiosClient";
 import { useToast } from "@/components/ui/use-toast";
 import { InfoCircledIcon } from "@radix-ui/react-icons";
 import {
@@ -34,6 +34,7 @@ import { ToastAction } from "@radix-ui/react-toast";
 import { Link } from "react-router-dom";
 import fetchToCurl from "fetch-to-curl";
 import { apiBaseUrl, envCredential } from "@/util/env";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 
 const createNewTaskFormSchema = z.object({
   url: z.string().url({
@@ -68,6 +69,7 @@ function createTaskRequestObject(formValues: CreateNewTaskFormValues) {
 function CreateNewTaskForm({ initialValues }: Props) {
   const queryClient = useQueryClient();
   const { toast } = useToast();
+  const credentialGetter = useCredentialGetter();
 
   const form = useForm<CreateNewTaskFormValues>({
     resolver: zodResolver(createNewTaskFormSchema),
@@ -75,8 +77,9 @@ function CreateNewTaskForm({ initialValues }: Props) {
   });
 
   const mutation = useMutation({
-    mutationFn: (formValues: CreateNewTaskFormValues) => {
+    mutationFn: async (formValues: CreateNewTaskFormValues) => {
       const taskRequest = createTaskRequestObject(formValues);
+      const client = await getClient(credentialGetter);
       return client.post<
         ReturnType<typeof createTaskRequestObject>,
         { data: { task_id: string } }
diff --git a/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx b/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx
index 4d7b0db6..8a62f6ab 100644
--- a/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx
@@ -1,4 +1,4 @@
-import { client } from "@/api/AxiosClient";
+import { getClient } from "@/api/AxiosClient";
 import {
   ArtifactApiResponse,
   ArtifactType,
@@ -16,6 +16,7 @@ import { TextArtifact } from "./TextArtifact";
 import { getImageURL } from "./artifactUtils";
 import { Input } from "@/components/ui/input";
 import { basicTimeFormat } from "@/util/timeFormat";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 
 type Props = {
   id: string;
@@ -24,6 +25,7 @@ type Props = {
 
 function StepArtifacts({ id, stepProps }: Props) {
   const { taskId } = useParams();
+  const credentialGetter = useCredentialGetter();
   const {
     data: artifacts,
     isFetching,
@@ -32,6 +34,7 @@ function StepArtifacts({ id, stepProps }: Props) {
   } = useQuery<Array<ArtifactApiResponse>>({
     queryKey: ["task", taskId, "steps", id, "artifacts"],
     queryFn: async () => {
+      const client = await getClient(credentialGetter);
       return client
         .get(`/tasks/${taskId}/steps/${id}/artifacts`)
         .then((response) => response.data);
diff --git a/skyvern-frontend/src/routes/tasks/detail/StepArtifactsLayout.tsx b/skyvern-frontend/src/routes/tasks/detail/StepArtifactsLayout.tsx
index 2d31f928..4b3e50c1 100644
--- a/skyvern-frontend/src/routes/tasks/detail/StepArtifactsLayout.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/StepArtifactsLayout.tsx
@@ -4,10 +4,12 @@ import { StepArtifacts } from "./StepArtifacts";
 import { useQuery } from "@tanstack/react-query";
 import { StepApiResponse } from "@/api/types";
 import { useParams } from "react-router-dom";
-import { client } from "@/api/AxiosClient";
+import { getClient } from "@/api/AxiosClient";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 
 function StepArtifactsLayout() {
   const [activeIndex, setActiveIndex] = useState(0);
+  const credentialGetter = useCredentialGetter();
   const { taskId } = useParams();
 
   const {
@@ -17,6 +19,7 @@ function StepArtifactsLayout() {
   } = useQuery<Array<StepApiResponse>>({
     queryKey: ["task", taskId, "steps"],
     queryFn: async () => {
+      const client = await getClient(credentialGetter);
       return client
         .get(`/tasks/${taskId}/steps`)
         .then((response) => response.data);
diff --git a/skyvern-frontend/src/routes/tasks/detail/StepNavigation.tsx b/skyvern-frontend/src/routes/tasks/detail/StepNavigation.tsx
index a1307fee..77b57955 100644
--- a/skyvern-frontend/src/routes/tasks/detail/StepNavigation.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/StepNavigation.tsx
@@ -1,10 +1,11 @@
-import { client } from "@/api/AxiosClient";
+import { getClient } from "@/api/AxiosClient";
 import { StepApiResponse } from "@/api/types";
 import { cn } from "@/util/utils";
 import { useQuery } from "@tanstack/react-query";
 import { useParams, useSearchParams } from "react-router-dom";
 import { PAGE_SIZE } from "../constants";
 import { CheckboxIcon, CrossCircledIcon } from "@radix-ui/react-icons";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 
 type Props = {
   activeIndex: number;
@@ -15,6 +16,7 @@ function StepNavigation({ activeIndex, onActiveIndexChange }: Props) {
   const { taskId } = useParams();
   const [searchParams] = useSearchParams();
   const page = searchParams.get("page") ? Number(searchParams.get("page")) : 1;
+  const credentialGetter = useCredentialGetter();
 
   const {
     data: steps,
@@ -23,6 +25,7 @@ function StepNavigation({ activeIndex, onActiveIndexChange }: Props) {
   } = useQuery<Array<StepApiResponse>>({
     queryKey: ["task", taskId, "steps", page],
     queryFn: async () => {
+      const client = await getClient(credentialGetter);
       return client
         .get(`/tasks/${taskId}/steps`, {
           params: {
diff --git a/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx b/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx
index ddca7c7b..63221a79 100644
--- a/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx
@@ -1,4 +1,4 @@
-import { client } from "@/api/AxiosClient";
+import { getClient } from "@/api/AxiosClient";
 import { Status, TaskApiResponse } from "@/api/types";
 import { Label } from "@/components/ui/label";
 import { Textarea } from "@/components/ui/textarea";
@@ -20,9 +20,11 @@ import {
 } from "@/components/ui/card";
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { Separator } from "@/components/ui/separator";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 
 function TaskDetails() {
   const { taskId } = useParams();
+  const credentialGetter = useCredentialGetter();
 
   const {
     data: task,
@@ -32,6 +34,7 @@ function TaskDetails() {
   } = useQuery<TaskApiResponse>({
     queryKey: ["task", taskId, "details"],
     queryFn: async () => {
+      const client = await getClient(credentialGetter);
       return client.get(`/tasks/${taskId}`).then((response) => response.data);
     },
     refetchInterval: (query) => {
diff --git a/skyvern-frontend/src/routes/tasks/list/TaskList.tsx b/skyvern-frontend/src/routes/tasks/list/TaskList.tsx
index 9d13e94b..e117c279 100644
--- a/skyvern-frontend/src/routes/tasks/list/TaskList.tsx
+++ b/skyvern-frontend/src/routes/tasks/list/TaskList.tsx
@@ -1,4 +1,4 @@
-import { client } from "@/api/AxiosClient";
+import { getClient } from "@/api/AxiosClient";
 import { TaskApiResponse } from "@/api/types";
 import { useQuery } from "@tanstack/react-query";
 import {
@@ -32,11 +32,13 @@ import {
   CardHeader,
   CardTitle,
 } from "@/components/ui/card";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 
 function TaskList() {
   const navigate = useNavigate();
   const [searchParams, setSearchParams] = useSearchParams();
   const page = searchParams.get("page") ? Number(searchParams.get("page")) : 1;
+  const credentialGetter = useCredentialGetter();
 
   const {
     data: tasks,
@@ -46,6 +48,7 @@ function TaskList() {
   } = useQuery<Array<TaskApiResponse>>({
     queryKey: ["tasks", "all", page],
     queryFn: async () => {
+      const client = await getClient(credentialGetter);
       const params = new URLSearchParams();
       params.append("page", String(page));
       params.append("page_size", String(PAGE_SIZE));
diff --git a/skyvern-frontend/src/routes/tasks/running/LatestScreenshot.tsx b/skyvern-frontend/src/routes/tasks/running/LatestScreenshot.tsx
index e81184af..f4f6dc63 100644
--- a/skyvern-frontend/src/routes/tasks/running/LatestScreenshot.tsx
+++ b/skyvern-frontend/src/routes/tasks/running/LatestScreenshot.tsx
@@ -1,4 +1,4 @@
-import { client } from "@/api/AxiosClient";
+import { getClient } from "@/api/AxiosClient";
 import {
   ArtifactApiResponse,
   ArtifactType,
@@ -7,12 +7,15 @@ import {
 import { Skeleton } from "@/components/ui/skeleton";
 import { keepPreviousData, useQuery } from "@tanstack/react-query";
 import { getImageURL } from "../detail/artifactUtils";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 
 type Props = {
   id: string;
 };
 
 function LatestScreenshot({ id }: Props) {
+  const credentialGetter = useCredentialGetter();
+
   const {
     data: artifact,
     isFetching,
@@ -20,6 +23,7 @@ function LatestScreenshot({ id }: Props) {
   } = useQuery<ArtifactApiResponse | undefined>({
     queryKey: ["task", id, "latestScreenshot"],
     queryFn: async () => {
+      const client = await getClient(credentialGetter);
       const steps: StepApiResponse[] = await client
         .get(`/tasks/${id}/steps`)
         .then((response) => response.data);
diff --git a/skyvern-frontend/src/routes/tasks/running/QueuedTasks.tsx b/skyvern-frontend/src/routes/tasks/running/QueuedTasks.tsx
index d93f0e8f..d92bb16d 100644
--- a/skyvern-frontend/src/routes/tasks/running/QueuedTasks.tsx
+++ b/skyvern-frontend/src/routes/tasks/running/QueuedTasks.tsx
@@ -1,4 +1,4 @@
-import { client } from "@/api/AxiosClient";
+import { getClient } from "@/api/AxiosClient";
 import { Status, TaskApiResponse } from "@/api/types";
 import { useQuery } from "@tanstack/react-query";
 import { basicTimeFormat } from "@/util/timeFormat";
@@ -12,12 +12,16 @@ import {
 } from "@/components/ui/table";
 import { useNavigate } from "react-router-dom";
 import { StatusBadge } from "@/components/StatusBadge";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 
 function QueuedTasks() {
   const navigate = useNavigate();
+  const credentialGetter = useCredentialGetter();
+
   const { data: tasks } = useQuery<Array<TaskApiResponse>>({
     queryKey: ["tasks", "queued"],
     queryFn: async () => {
+      const client = await getClient(credentialGetter);
       return client
         .get("/tasks", {
           params: {
diff --git a/skyvern-frontend/src/routes/tasks/running/RunningTasks.tsx b/skyvern-frontend/src/routes/tasks/running/RunningTasks.tsx
index 7261bd99..1e86c6af 100644
--- a/skyvern-frontend/src/routes/tasks/running/RunningTasks.tsx
+++ b/skyvern-frontend/src/routes/tasks/running/RunningTasks.tsx
@@ -1,4 +1,4 @@
-import { client } from "@/api/AxiosClient";
+import { getClient } from "@/api/AxiosClient";
 import { TaskApiResponse } from "@/api/types";
 import { useQuery } from "@tanstack/react-query";
 import { useNavigate } from "react-router-dom";
@@ -12,13 +12,16 @@ import {
 } from "@/components/ui/card";
 import { basicTimeFormat } from "@/util/timeFormat";
 import { LatestScreenshot } from "./LatestScreenshot";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 
 function RunningTasks() {
   const navigate = useNavigate();
+  const credentialGetter = useCredentialGetter();
 
   const { data: runningTasks } = useQuery<Array<TaskApiResponse>>({
     queryKey: ["tasks", "running"],
     queryFn: async () => {
+      const client = await getClient(credentialGetter);
       return client
         .get("/tasks", {
           params: {
diff --git a/skyvern-frontend/src/store/CredentialGetterContext.ts b/skyvern-frontend/src/store/CredentialGetterContext.ts
new file mode 100644
index 00000000..f5f7d78a
--- /dev/null
+++ b/skyvern-frontend/src/store/CredentialGetterContext.ts
@@ -0,0 +1,6 @@
+import { CredentialGetter } from "@/api/AxiosClient";
+import { createContext } from "react";
+
+const CredentialGetterContext = createContext<CredentialGetter | null>(null);
+
+export { CredentialGetterContext };