fix(frontend): socket logout handling (#10445)

## Changes 🏗️

- Close websocket connections gracefully during logout ( _whether from
another tab or not_ )
- Also fixed an error on `HeroSection` that shows when the onboarding is
disabled locally ( `null` )
- Uncomment legit tests about connecting/saving agents
- Centralise local storage usage through a single service with typed
keys

## Checklist 📋

### For code changes:
- [x] I have clearly listed my changes in the PR description
- [x] I have made a test plan
- [x] I have tested my changes according to the test plan:
  - [x] Login in 3 tabs ( 1 builder, 1 marketplace, 1 agent run view )
  - [x] Logout from the marketplace tab
- [x] The other tabs show logout state gracefully without toasts or
errors
- [x] Websocket connections are closed ( _devtools console shows that_ )

### For configuration changes:

None

autogpt_platform/frontend/src/app/(platform)/library/agents/[id]/page.tsx

@@ -227,8 +227,8 @@ export default function AgentRunsPage(): React.ReactElement {
             <LoadingSpinner className="ml-1.5 size-3.5" />
           </div>
         ),
-        duration: Infinity, // show until connection is re-established
-        dismissable: false,
+        duration: Infinity,
+        dismissable: true,
       });
     });
     const cancelConnectHandler = api.onWebSocketConnect(() => {

autogpt_platform/frontend/src/components/Flow.tsx

@@ -36,6 +36,7 @@ import {
   LibraryAgent,
 } from "@/lib/autogpt-server-api";
 import { useBackendAPI } from "@/lib/autogpt-server-api/context";
+import { Key, storage } from "@/services/storage/local-storage";
 import { getTypeColor, findNewlyAddedBlockCoordinates } from "@/lib/utils";
 import { history } from "./history";
 import { CustomEdge } from "./CustomEdge";
@@ -157,8 +158,6 @@ const FlowEditor: React.FC<{
 
   const { toast } = useToast();
 
-  const TUTORIAL_STORAGE_KEY = "shepherd-tour";
-
   // It stores the dimension of all nodes with position as well
   const [nodeDimensions, setNodeDimensions] = useState<NodeDimension>({});
 
@@ -181,13 +180,13 @@ const FlowEditor: React.FC<{
 
   useEffect(() => {
     if (params.get("resetTutorial") === "true") {
-      localStorage.removeItem(TUTORIAL_STORAGE_KEY);
+      storage.clean(Key.SHEPHERD_TOUR);
       router.push(pathname);
-    } else if (!localStorage.getItem(TUTORIAL_STORAGE_KEY)) {
+    } else if (!storage.get(Key.SHEPHERD_TOUR)) {
       const emptyNodes = (forceRemove: boolean = false) =>
         forceRemove ? (setNodes([]), setEdges([]), true) : nodes.length === 0;
       startTutorial(emptyNodes, setPinBlocksPopover, setPinSavePopover);
-      localStorage.setItem(TUTORIAL_STORAGE_KEY, "yes");
+      storage.set(Key.SHEPHERD_TOUR, "yes");
     }
   }, [router, pathname, params, setEdges, setNodes, nodes.length]);
 

autogpt_platform/frontend/src/components/agptui/Wallet.tsx

@@ -20,12 +20,15 @@ export default function Wallet() {
   const { credits, formatCredits, fetchCredits } = useCredits({
     fetchInitialCredits: true,
   });
+
   const { state, updateState } = useOnboarding();
   const [prevCredits, setPrevCredits] = useState<number | null>(credits);
   const [flash, setFlash] = useState(false);
+
   const [stepsLength, setStepsLength] = useState<number | null>(
     state?.completedSteps?.length || null,
   );
+
   const walletRef = useRef<HTMLButtonElement | null>(null);
 
   const onWalletOpen = useCallback(async () => {

autogpt_platform/frontend/src/components/agptui/composite/HeroSection.tsx

@@ -8,12 +8,12 @@ import { useOnboarding } from "@/components/onboarding/onboarding-provider";
 
 export const HeroSection: React.FC = () => {
   const router = useRouter();
-  const { completeStep } = useOnboarding();
+  const onboarding = useOnboarding();
 
   // Mark marketplace visit task as completed
   React.useEffect(() => {
-    completeStep("MARKETPLACE_VISIT");
-  }, [completeStep]);
+    onboarding?.completeStep("MARKETPLACE_VISIT");
+  }, [onboarding?.completeStep]);
 
   function onFilterChange(selectedFilters: string[]) {
     const encodedTerm = encodeURIComponent(selectedFilters.join(", "));

autogpt_platform/frontend/src/components/tutorial.ts

@@ -1,6 +1,7 @@
 import Shepherd from "shepherd.js";
 import "shepherd.js/dist/css/shepherd.css";
 import { sendGAEvent } from "@/components/analytics/google-analytics";
+import { Key, storage } from "@/services/storage/local-storage";
 
 export const startTutorial = (
   emptyNodeList: (forceEmpty: boolean) => boolean,
@@ -152,7 +153,7 @@ export const startTutorial = (
         text: "Skip Tutorial",
         action: () => {
           tour.cancel(); // Ends the tour
-          localStorage.setItem("shepherd-tour", "skipped"); // Set the tutorial as skipped in local storage
+          storage.set(Key.SHEPHERD_TOUR, "skipped"); // Set the tutorial as skipped in local storage
         },
         classes: "shepherd-button-secondary", // Optionally add a class for styling the skip button differently
       },
@@ -546,7 +547,7 @@ export const startTutorial = (
   tour.on("complete", () => {
     setPinBlocksPopover(false);
     setPinSavePopover(false);
-    localStorage.setItem("shepherd-tour", "completed"); // Optionally mark the tutorial as completed
+    storage.set(Key.SHEPHERD_TOUR, "completed"); // Optionally mark the tutorial as completed
   });
 
   for (const step of tour.steps) {
@@ -561,7 +562,7 @@ export const startTutorial = (
   tour.on("cancel", () => {
     setPinBlocksPopover(false);
     setPinSavePopover(false);
-    localStorage.setItem("shepherd-tour", "canceled"); // Optionally mark the tutorial as canceled
+    storage.set(Key.SHEPHERD_TOUR, "canceled"); // Optionally mark the tutorial as canceled
   });
 
   tour.start();

autogpt_platform/frontend/src/hooks/useCopyPaste.ts

@@ -1,5 +1,6 @@
 import { useCallback } from "react";
 import { Node, Edge, useReactFlow } from "@xyflow/react";
+import { Key, storage } from "@/services/storage/local-storage";
 
 interface CopyableData {
   nodes: Node[];
@@ -37,10 +38,10 @@ export function useCopyPaste(getNextNodeId: () => string) {
             edges: selectedEdges,
           };
 
-          localStorage.setItem("copiedFlowData", JSON.stringify(copiedData));
+          storage.set(Key.COPIED_FLOW_DATA, JSON.stringify(copiedData));
         }
         if (event.key === "v" || event.key === "V") {
-          const copiedDataString = localStorage.getItem("copiedFlowData");
+          const copiedDataString = storage.get(Key.COPIED_FLOW_DATA);
           if (copiedDataString) {
             const copiedData = JSON.parse(copiedDataString) as CopyableData;
             const oldToNewIdMap: Record<string, string> = {};

autogpt_platform/frontend/src/lib/autogpt-server-api/client.ts

@@ -2,6 +2,8 @@ import { getWebSocketToken } from "@/lib/supabase/actions";
 import { getServerSupabase } from "@/lib/supabase/server/getServerSupabase";
 import { createBrowserClient } from "@supabase/ssr";
 import type { SupabaseClient } from "@supabase/supabase-js";
+import { Key, storage } from "@/services/storage/local-storage";
+import * as Sentry from "@sentry/nextjs";
 import type {
   AddUserCreditsResponse,
   AnalyticsDetails,
@@ -1141,6 +1143,7 @@ export default class BackendAPI {
           this.webSocket!.state = "connected";
           console.info("[BackendAPI] WebSocket connected to", this.wsUrl);
           this._startWSHeartbeat(); // Start heartbeat when connection opens
+          this._clearDisconnectIntent(); // Clear disconnect intent when connected
           this.wsOnConnectHandlers.forEach((handler) => handler());
           resolve();
         };
@@ -1161,10 +1164,12 @@ export default class BackendAPI {
 
           this._stopWSHeartbeat(); // Stop heartbeat when connection closes
           this.wsConnecting = null;
-          this.wsOnDisconnectHandlers.forEach((handler) => handler());
 
-          // Only attempt to reconnect if this wasn't an intentional disconnection
-          if (!this.isIntentionallyDisconnected) {
+          const wasIntentional =
+            this.isIntentionallyDisconnected || this._hasDisconnectIntent();
+
+          if (!wasIntentional) {
+            this.wsOnDisconnectHandlers.forEach((handler) => handler());
             setTimeout(() => this.connectWebSocket().then(resolve), 1000);
           }
         };
@@ -1174,7 +1179,6 @@ export default class BackendAPI {
             console.error("[BackendAPI] WebSocket error:", error);
           }
         };
-
         this.webSocket.onmessage = (event) => this._handleWSMessage(event);
       } catch (error) {
         console.error("[BackendAPI] Error connecting to WebSocket:", error);
@@ -1191,6 +1195,28 @@ export default class BackendAPI {
     }
   }
 
+  private _hasDisconnectIntent(): boolean {
+    if (!isClient) return false;
+
+    try {
+      return storage.get(Key.WEBSOCKET_DISCONNECT_INTENT) === "true";
+    } catch {
+      return false;
+    }
+  }
+
+  private _clearDisconnectIntent(): void {
+    if (!isClient) return;
+
+    try {
+      storage.clean(Key.WEBSOCKET_DISCONNECT_INTENT);
+    } catch {
+      Sentry.captureException(
+        new Error("Failed to clear WebSocket disconnect intent"),
+      );
+    }
+  }
+
   private _handleWSMessage(event: MessageEvent): void {
     const message: WebsocketMessage = JSON.parse(event.data);
 

autogpt_platform/frontend/src/lib/autogpt-server-api/helpers.ts

@@ -1,4 +1,5 @@
 import { getServerSupabase } from "@/lib/supabase/server/getServerSupabase";
+import { Key, storage } from "@/services/storage/local-storage";
 
 export class ApiError extends Error {
   public status: number;
@@ -182,7 +183,7 @@ function isLogoutInProgress(): boolean {
 
   try {
     // Check if logout was recently triggered
-    const logoutTimestamp = window.localStorage.getItem("supabase-logout");
+    const logoutTimestamp = storage.get(Key.LOGOUT);
     if (logoutTimestamp) {
       const timeDiff = Date.now() - parseInt(logoutTimestamp);
       // Consider logout in progress for 5 seconds after trigger

autogpt_platform/frontend/src/lib/supabase/helpers.ts

@@ -1,5 +1,6 @@
 import { type CookieOptions } from "@supabase/ssr";
 import { SupabaseClient } from "@supabase/supabase-js";
+import { Key, storage } from "@/services/storage/local-storage";
 
 export const PROTECTED_PAGES = [
   "/monitor",
@@ -12,10 +13,6 @@ export const PROTECTED_PAGES = [
 
 export const ADMIN_PAGES = ["/admin"] as const;
 
-export const STORAGE_KEYS = {
-  LOGOUT: "supabase-logout",
-} as const;
-
 export function getCookieSettings(): Partial<CookieOptions> {
   return {
     secure: process.env.NODE_ENV === "production",
@@ -39,13 +36,24 @@ export function shouldRedirectOnLogout(pathname: string): boolean {
 
 // Cross-tab logout utilities
 export function broadcastLogout(): void {
-  if (typeof window !== "undefined") {
-    window.localStorage.setItem(STORAGE_KEYS.LOGOUT, Date.now().toString());
-  }
+  storage.set(Key.LOGOUT, Date.now().toString());
 }
 
 export function isLogoutEvent(event: StorageEvent): boolean {
-  return event.key === STORAGE_KEYS.LOGOUT;
+  return event.key === Key.LOGOUT;
+}
+
+// WebSocket disconnect intent utilities
+export function setWebSocketDisconnectIntent(): void {
+  storage.set(Key.WEBSOCKET_DISCONNECT_INTENT, "true");
+}
+
+export function clearWebSocketDisconnectIntent(): void {
+  storage.clean(Key.WEBSOCKET_DISCONNECT_INTENT);
+}
+
+export function hasWebSocketDisconnectIntent(): boolean {
+  return storage.get(Key.WEBSOCKET_DISCONNECT_INTENT) === "true";
 }
 
 // Redirect utilities

autogpt_platform/frontend/src/lib/supabase/hooks/useSupabase.ts

@@ -13,8 +13,10 @@ import {
 } from "../actions";
 import {
   broadcastLogout,
+  clearWebSocketDisconnectIntent,
   getRedirectPath,
   isLogoutEvent,
+  setWebSocketDisconnectIntent,
   setupSessionEventListeners,
 } from "../helpers";
 
@@ -47,6 +49,7 @@ export function useSupabase() {
   }, []);
 
   async function logOut(options: ServerLogoutOptions = {}) {
+    setWebSocketDisconnectIntent();
     api.disconnectWebSocket();
     broadcastLogout();
 
@@ -93,6 +96,7 @@ export function useSupabase() {
           }
           return currentUser;
         });
+        clearWebSocketDisconnectIntent();
       }
 
       return true;
@@ -119,6 +123,7 @@ export function useSupabase() {
       }
 
       setUser(serverUser);
+      clearWebSocketDisconnectIntent();
       return serverUser;
     } catch (error) {
       console.error("Get user error:", error);
@@ -130,6 +135,7 @@ export function useSupabase() {
   function handleCrossTabLogout(e: StorageEvent) {
     if (!isLogoutEvent(e)) return;
 
+    setWebSocketDisconnectIntent();
     api.disconnectWebSocket();
 
     // Clear local state immediately

autogpt_platform/frontend/src/services/storage/local-storage.ts

@@ -0,0 +1,43 @@
+import * as Sentry from "@sentry/nextjs";
+
+export enum Key {
+  LOGOUT = "supabase-logout",
+  WEBSOCKET_DISCONNECT_INTENT = "websocket-disconnect-intent",
+  COPIED_FLOW_DATA = "copied-flow-data",
+  SHEPHERD_TOUR = "shepherd-tour",
+}
+
+function get(key: Key) {
+  if (typeof window === "undefined") {
+    Sentry.captureException(new Error("Local storage is not available"));
+    return;
+  }
+  try {
+    return window.localStorage.getItem(key);
+  } catch {
+    // Fine, just return undefined not always items will be set on local storage
+    return;
+  }
+}
+
+function set(key: Key, value: string) {
+  if (typeof window === "undefined") {
+    Sentry.captureException(new Error("Local storage is not available"));
+    return;
+  }
+  return window.localStorage.setItem(key, value);
+}
+
+function clean(key: Key) {
+  if (typeof window === "undefined") {
+    Sentry.captureException(new Error("Local storage is not available"));
+    return;
+  }
+  return window.localStorage.removeItem(key);
+}
+
+export const storage = {
+  clean,
+  get,
+  set,
+};

autogpt_platform/frontend/src/tests/build.spec.ts

@@ -257,16 +257,11 @@ test.describe("Build", () => { //(1)!
     // Ensure the run button is enabled
     await test.expect(buildPage.isRunButtonEnabled()).resolves.toBeTruthy();
 
-    // Run the agent
-    // await buildPage.runAgent();
-
-    // Wait for processing to complete by checking the completion badge
-    // await buildPage.waitForCompletionBadge();
-
-    // Get the first completion badge and verify it's visible
-    // await test
-    //   .expect(buildPage.isCompletionBadgeVisible())
-    //   .resolves.toBeTruthy();
+    await buildPage.runAgent();
+    await buildPage.waitForCompletionBadge();
+    await test
+      .expect(buildPage.isCompletionBadgeVisible())
+      .resolves.toBeTruthy();
   });
 
   test("user can build an agent with inputs and output blocks", async ({ page }, testInfo) => {
@@ -355,15 +350,14 @@ test.describe("Build", () => { //(1)!
     // Wait for save to complete
     await page.waitForTimeout(1000);
 
-    // TODO: investigate why running 
-    // await buildPage.runAgent();
-    // await buildPage.fillRunDialog({
-    //   Value: "10",
-    // });
-    // await buildPage.clickRunDialogRunButton();
-    // await buildPage.waitForCompletionBadge();
-    // await test
-    //   .expect(buildPage.isCompletionBadgeVisible())
-    //   .resolves.toBeTruthy();
+    await buildPage.runAgent();
+    await buildPage.fillRunDialog({
+      Value: "10",
+    });
+    await buildPage.clickRunDialogRunButton();
+    await buildPage.waitForCompletionBadge();
+    await test
+      .expect(buildPage.isCompletionBadgeVisible())
+      .resolves.toBeTruthy();
   });
 });