Jon/keyed browser session id cache (#2928)

skyvern-frontend/src/routes/workflows/editor/nodes/components/NodeHeader.tsx

@@ -5,7 +5,7 @@ import { useLocation, useNavigate, useParams } from "react-router-dom";
 import { useMutation, useQueryClient } from "@tanstack/react-query";
 
 import { getClient } from "@/api/AxiosClient";
-import { ProxyLocation } from "@/api/types";
+import { ProxyLocation, User } from "@/api/types";
 import { Timer } from "@/components/Timer";
 import { toast } from "@/components/ui/use-toast";
 import { useCredentialGetter } from "@/hooks/useCredentialGetter";
@@ -29,11 +29,15 @@ import {
   statusIsFinalized,
   statusIsRunningOrQueued,
 } from "@/routes/tasks/types";
+import {
+  useOptimisticallyRequestBrowserSessionId,
+  type OptimisticBrowserSession,
+} from "@/store/useOptimisticallyRequestBrowserSessionId";
+import { useUser } from "@/hooks/useUser";
 
 import { EditableNodeTitle } from "../components/EditableNodeTitle";
 import { NodeActionMenu } from "../NodeActionMenu";
 import { WorkflowBlockIcon } from "../WorkflowBlockIcon";
-import { lsKeys } from "@/util/env";
 
 interface Props {
   blockLabel: string; // today, this + wpid act as the identity of a block
@@ -69,12 +73,24 @@ const blockTypeToTitle = (type: WorkflowBlockType): string => {
 
 const getPayload = (opts: {
   blockLabel: string;
+  optimistic: OptimisticBrowserSession;
   parameters: Record<string, unknown>;
   totpIdentifier: string | null;
   totpUrl: string | null;
+  user: User | null;
   workflowPermanentId: string;
   workflowSettings: WorkflowSettingsState;
-}): Payload => {
+}): Payload | null => {
+  if (!opts.user) {
+    toast({
+      variant: "warning",
+      title: "Error",
+      description: "No user found",
+    });
+
+    return null;
+  }
+
   const webhook_url = opts.workflowSettings.webhookCallbackUrl.trim();
 
   let extraHttpHeaders = null;
@@ -92,15 +108,12 @@ const getPayload = (opts: {
     });
   }
 
-  const stored = localStorage.getItem(lsKeys.optimisticBrowserSession);
-  let browserSessionId: string | null = null;
-  try {
-    const parsed = JSON.parse(stored ?? "");
-    const { browser_session_id } = parsed;
-    browserSessionId = browser_session_id as string;
-  } catch {
-    // pass
-  }
+  const browserSessionData = opts.optimistic.get(
+    opts.user,
+    opts.workflowPermanentId,
+  );
+
+  const browserSessionId = browserSessionData?.browser_session_id;
 
   if (!browserSessionId) {
     toast({
@@ -108,6 +121,8 @@ const getPayload = (opts: {
       title: "Error",
       description: "No browser session ID found",
     });
+
+    return null;
   } else {
     toast({
       variant: "default",
@@ -166,6 +181,8 @@ function NodeHeader({
   const { data: workflowRun } = useWorkflowRunQuery();
   const workflowRunIsRunningOrQueued =
     workflowRun && statusIsRunningOrQueued(workflowRun);
+  const optimistic = useOptimisticallyRequestBrowserSessionId();
+  const user = useUser().get();
 
   useEffect(() => {
     if (!workflowRun || !workflowPermanentId || !workflowRunId) {
@@ -227,13 +244,19 @@ function NodeHeader({
 
       const body = getPayload({
         blockLabel,
+        optimistic,
         parameters,
         totpIdentifier,
         totpUrl,
+        user,
         workflowPermanentId,
         workflowSettings: workflowSettingsStore,
       });
 
+      if (!body) {
+        return;
+      }
+
       return await client.post<Payload, { data: { run_id: string } }>(
         "/run/workflows/blocks",
         body,
@@ -268,8 +291,10 @@ function NodeHeader({
   const cancelBlock = useMutation({
     mutationFn: async () => {
       const browserSessionId =
-        debugStore.getCurrentBrowserSessionId() ??
-        "<missing-browser-session-id>";
+        user && workflowPermanentId
+          ? optimistic.get(user, workflowPermanentId)?.browser_session_id ??
+            "<missing-browser-session-id>"
+          : "<missing-user-or-workflow-permanent-id>";
       const client = await getClient(credentialGetter);
       return client
         .post(`/runs/${browserSessionId}/workflow_run/${workflowRunId}/cancel/`)

skyvern-frontend/src/store/DebugStoreContext.tsx

@@ -1,6 +1,5 @@
 import React, { createContext, useMemo } from "react";
 import { useLocation } from "react-router-dom";
-import { lsKeys } from "@/util/env";
 
 function useIsDebugMode() {
   const location = useLocation();
@@ -10,23 +9,8 @@ function useIsDebugMode() {
   );
 }
 
-function getCurrentBrowserSessionId() {
-  const stored = localStorage.getItem(lsKeys.optimisticBrowserSession);
-  let browserSessionId: string | null = null;
-  try {
-    const parsed = JSON.parse(stored ?? "");
-    const { browser_session_id } = parsed;
-    browserSessionId = browser_session_id as string;
-  } catch {
-    // pass
-  }
-
-  return browserSessionId;
-}
-
 export type DebugStoreContextType = {
   isDebugMode: boolean;
-  getCurrentBrowserSessionId: () => string | null;
 };
 
 export const DebugStoreContext = createContext<
@@ -39,9 +23,7 @@ export const DebugStoreProvider: React.FC<{ children: React.ReactNode }> = ({
   const isDebugMode = useIsDebugMode();
 
   return (
-    <DebugStoreContext.Provider
-      value={{ isDebugMode, getCurrentBrowserSessionId }}
-    >
+    <DebugStoreContext.Provider value={{ isDebugMode }}>
       {children}
     </DebugStoreContext.Provider>
   );

skyvern-frontend/src/store/useOptimisticallyRequestBrowserSessionId.ts

@@ -1,5 +1,7 @@
-import { create } from "zustand";
 import { AxiosInstance } from "axios";
+import { create as createStore } from "zustand";
+
+import { User } from "@/api/types";
 import { lsKeys } from "@/util/env";
 
 export interface BrowserSessionData {
@@ -7,59 +9,120 @@ export interface BrowserSessionData {
   expires_at: number | null; // seconds since epoch
 }
 
-interface OptimisticBrowserSessionIdState extends BrowserSessionData {
-  run: (client: AxiosInstance) => Promise<BrowserSessionData>;
+interface RunOpts {
+  client: AxiosInstance;
+  reason?: string;
+  user: User;
+  workflowPermanentId?: string;
+}
+
+export interface OptimisticBrowserSession {
+  get: (user: User, workflowPermanentId: string) => BrowserSessionData | null;
+  run: (runOpts: RunOpts) => Promise<BrowserSessionData>;
 }
 
 const SESSION_TIMEOUT_MINUTES = 60;
+const SPARE = "spare";
+
+const makeKey = (user: User, workflowPermanentId?: string | undefined) => {
+  return `${lsKeys.optimisticBrowserSession}:${user.id}:${workflowPermanentId ?? SPARE}`;
+};
+
+/**
+ * Read a `BrowserSessionData` from localStorage cache. If the entry is expired,
+ * return `null`. If the entry is invalid, return `null`. Otherwise return it.
+ */
+const read = (key: string): BrowserSessionData | null => {
+  const stored = localStorage.getItem(key);
+  if (stored) {
+    try {
+      const parsed = JSON.parse(stored);
+      const { browser_session_id, expires_at } = parsed;
+      const now = Math.floor(Date.now() / 1000); // seconds since epoch
+
+      if (
+        browser_session_id &&
+        typeof browser_session_id === "string" &&
+        expires_at &&
+        typeof expires_at === "number" &&
+        now < expires_at
+      ) {
+        return { browser_session_id, expires_at };
+      }
+    } catch (e) {
+      // pass
+    }
+  }
+
+  return null;
+};
+
+/**
+ * Write a `BrowserSessionData` to localStorage cache.
+ */
+const write = (key: string, browserSessionData: BrowserSessionData) => {
+  localStorage.setItem(key, JSON.stringify(browserSessionData));
+};
+
+/**
+ * Delete a localStorage key.
+ */
+const del = (key: string) => {
+  localStorage.removeItem(key);
+};
+
+/**
+ * Create a new browser session and return the `BrowserSessionData`.
+ */
+const create = async (client: AxiosInstance): Promise<BrowserSessionData> => {
+  const resp = await client.post("/browser_sessions", {
+    timeout: SESSION_TIMEOUT_MINUTES,
+  });
+
+  const { browser_session_id: newBrowserSessionId, timeout } = resp.data;
+  const newExpiresAt = Math.floor(Date.now() / 1000) + timeout * 60 * 0.9;
+
+  return {
+    browser_session_id: newBrowserSessionId,
+    expires_at: newExpiresAt,
+  };
+};
 
 export const useOptimisticallyRequestBrowserSessionId =
-  create<OptimisticBrowserSessionIdState>((set) => ({
-    browser_session_id: null,
-    expires_at: null,
-    run: async (client) => {
-      const stored = localStorage.getItem(lsKeys.optimisticBrowserSession);
-      if (stored) {
-        try {
-          const parsed = JSON.parse(stored);
-          const { browser_session_id, expires_at } = parsed;
-          const now = Math.floor(Date.now() / 1000); // seconds since epoch
+  createStore<OptimisticBrowserSession>(() => ({
+    get: (user: User, workflowPermanentId: string) => {
+      return read(makeKey(user, workflowPermanentId));
+    },
+    run: async ({ client, user, workflowPermanentId }: RunOpts) => {
+      if (workflowPermanentId) {
+        const userKey = makeKey(user, workflowPermanentId);
+        const exists = read(userKey);
 
-          if (
-            browser_session_id &&
-            typeof browser_session_id === "string" &&
-            expires_at &&
-            typeof expires_at === "number" &&
-            now < expires_at
-          ) {
-            set({ browser_session_id, expires_at });
-            return { browser_session_id, expires_at };
-          }
-        } catch (e) {
-          // pass
+        if (exists) {
+          return exists;
+        }
+
+        const spareKey = makeKey(user, SPARE);
+        const spare = read(spareKey);
+
+        if (spare) {
+          del(spareKey);
+          write(userKey, spare);
+          create(client).then((newSpare) => write(spareKey, newSpare));
+          return spare;
         }
       }
 
-      const resp = await client.post("/browser_sessions", {
-        timeout: SESSION_TIMEOUT_MINUTES,
-      });
-      const { browser_session_id: newBrowserSessionId, timeout } = resp.data;
-      const newExpiresAt = Math.floor(Date.now() / 1000) + timeout * 60 * 0.9;
-      set({
-        browser_session_id: newBrowserSessionId,
-        expires_at: newExpiresAt,
-      });
-      localStorage.setItem(
-        lsKeys.optimisticBrowserSession,
-        JSON.stringify({
-          browser_session_id: newBrowserSessionId,
-          expires_at: newExpiresAt,
-        }),
-      );
+      const key = makeKey(user, workflowPermanentId);
+      const browserSessionData = read(key);
 
-      return {
-        browser_session_id: newBrowserSessionId,
-        expires_at: newExpiresAt,
-      };
+      if (browserSessionData) {
+        return browserSessionData;
+      }
+
+      const knew = await create(client);
+      write(key, knew);
+
+      return knew;
     },
   }));