fix: Enhance state management with reset methods and logout functionality improvements (#7234)

* feat: add resetStore method to foldersStore for state management

✨ (foldersStore.tsx): Implement resetStore function to clear folders and reset related state variables, enhancing state management capabilities.

* feat: add resetStore method to flowsManagerStore for improved state management

✨ (flowsManagerStore.ts): Implement resetStore function to clear flows and reset related state variables, enhancing the store's state management capabilities.

* feat: add resetFlowState method to flowStore for enhanced state management

✨ (flowStore.ts): Implement resetFlowState function to clear and reset the flow-related state variables, improving the management of flow state in the application.

* feat: enhance logout functionality with state resets and query invalidation

✨ (use-post-logout.ts): Update the useLogout function to reset flow, flows manager, and folder states upon successful logout. Additionally, invalidate and remove related queries to ensure data consistency in the application.

* test: add user flow state cleanup test for session management

✨ (user-flow-state-cleanup.spec.ts): Implement a comprehensive test to verify that user flow states are properly cleaned up between sessions. The test includes user creation, flow management, and ensures that the admin cannot see user-specific flows after logout, enhancing the robustness of session handling in the application.

* refactor: streamline logout query management in useLogout function

🔧 (use-post-logout.ts): Remove redundant query removal calls in the useLogout function, focusing on invalidating queries to maintain data consistency post-logout. This change simplifies the logout process and enhances code clarity.

* 🔧 (user-flow-state-cleanup.spec.ts): remove unnecessary verification step for admin empty state to improve test readability and efficiency

* ✨ (user-flow-state-cleanup.spec.ts): improve user flow state cleanup test by adding error handling for clicking on new project button and new flow button to ensure test stability and reliability

---------

Co-authored-by: cristhianzl <cristhian.lousa@gmail.com>

src/frontend/package-lock.json

@@ -706,7 +706,6 @@
     },
     "node_modules/@clack/prompts/node_modules/is-unicode-supported": {
       "version": "1.3.0",
-      "extraneous": true,
       "inBundle": true,
       "license": "MIT",
       "engines": {

src/frontend/src/controllers/API/queries/auth/use-post-logout.ts

@@ -5,6 +5,9 @@ import {
   IS_AUTO_LOGIN,
   LANGFLOW_AUTO_LOGIN_OPTION,
 } from "@/constants/constants";
+import useFlowStore from "@/stores/flowStore";
+import useFlowsManagerStore from "@/stores/flowsManagerStore";
+import { useFolderStore } from "@/stores/foldersStore";
 import { Cookies } from "react-cookie";
 import { api } from "../../api";
 import { getURL } from "../../helpers/constants";
@@ -13,7 +16,7 @@ import { UseRequestProcessor } from "../../services/request-processor";
 export const useLogout: useMutationFunctionType<undefined, void> = (
   options?,
 ) => {
-  const { mutate } = UseRequestProcessor();
+  const { mutate, queryClient } = UseRequestProcessor();
   const cookies = new Cookies();
   const logout = useAuthStore((state) => state.logout);
   const isAutoLoginEnv = IS_AUTO_LOGIN;
@@ -34,6 +37,14 @@ export const useLogout: useMutationFunctionType<undefined, void> = (
   const mutation = mutate(["useLogout"], logoutUser, {
     onSuccess: () => {
       logout();
+
+      useFlowStore.getState().resetFlowState();
+      useFlowsManagerStore.getState().resetStore();
+      useFolderStore.getState().resetStore();
+
+      queryClient.invalidateQueries({ queryKey: ["useGetRefreshFlowsQuery"] });
+      queryClient.invalidateQueries({ queryKey: ["useGetFolders"] });
+      queryClient.invalidateQueries({ queryKey: ["useGetFolder"] });
     },
     onError: (error) => {
       console.error(error);

src/frontend/src/stores/flowStore.ts

@@ -970,6 +970,26 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
   setCurrentBuildingNodeId: (nodeIds) => {
     set({ currentBuildingNodeId: nodeIds });
   },
+  resetFlowState: () => {
+    set({
+      nodes: [],
+      edges: [],
+      flowState: undefined,
+      hasIO: false,
+      inputs: [],
+      outputs: [],
+      flowPool: {},
+      currentFlow: undefined,
+      reactFlowInstance: null,
+      lastCopiedSelection: null,
+      verticesBuild: null,
+      flowBuildStatus: {},
+      isBuilding: false,
+      isPending: true,
+      positionDictionary: {},
+      componentsToUpdate: [],
+    });
+  },
 }));
 
 export default useFlowStore;

src/frontend/src/stores/flowsManagerStore.ts

@@ -137,6 +137,16 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
       resolve();
     });
   },
+  resetStore: () => {
+    set({
+      flows: [],
+      currentFlow: undefined,
+      currentFlowId: "",
+      flowToCanvas: null,
+      searchFlowsComponents: "",
+      selectedFlowsComponentsCards: [],
+    });
+  },
 }));
 
 export default useFlowsManagerStore;

src/frontend/src/stores/foldersStore.tsx

@@ -17,4 +17,13 @@ export const useFolderStore = create<FoldersStoreType>((set, get) => ({
   setStarterProjectId: (id) => set(() => ({ starterProjectId: id })),
   folders: [],
   setFolders: (folders) => set(() => ({ folders: folders })),
+  resetStore: () => {
+    set({
+      folders: [],
+      myCollectionId: "",
+      folderToEdit: null,
+      folderDragging: false,
+      folderIdDragging: "",
+    });
+  },
 }));

src/frontend/src/types/zustand/flow/index.ts

@@ -95,6 +95,7 @@ export type FlowStoreType = {
   setIsBuilding: (isBuilding: boolean) => void;
   setPending: (isPending: boolean) => void;
   resetFlow: (flow: FlowType | undefined) => void;
+  resetFlowState: () => void;
   reactFlowInstance: ReactFlowInstance<AllNodeType, EdgeType> | null;
   setReactFlowInstance: (
     newState: ReactFlowInstance<AllNodeType, EdgeType>,

src/frontend/src/types/zustand/flowsManager/index.ts

@@ -30,6 +30,7 @@ export type FlowsManagerStoreType = {
   setFlowToCanvas: (flowToCanvas: FlowType | null) => Promise<void>;
   IOModalOpen: boolean;
   setIOModalOpen: (IOModalOpen: boolean) => void;
+  resetStore: () => void;
 };
 
 export type UseUndoRedoOptions = {

src/frontend/src/types/zustand/folders/index.ts

@@ -13,4 +13,5 @@ export type FoldersStoreType = {
   setStarterProjectId: (id: string) => void;
   folders: FolderType[];
   setFolders: (folders: FolderType[]) => void;
+  resetStore: () => void;
 };

src/frontend/tests/core/features/user-flow-state-cleanup.spec.ts

@@ -0,0 +1,160 @@
+import { expect, test } from "@playwright/test";
+import { awaitBootstrapTest } from "../../utils/await-bootstrap-test";
+
+test(
+  "flow state should be properly cleaned up between user sessions",
+  { tag: ["@bug-fix", "@api", "@database"] },
+  async ({ page }) => {
+    // Disable auto login
+    await page.route("**/api/v1/auto_login", (route) => {
+      route.fulfill({
+        status: 500,
+        contentType: "application/json",
+        body: JSON.stringify({
+          detail: { auto_login: false },
+        }),
+      });
+    });
+
+    await page.addInitScript(() => {
+      window.process = window.process || {};
+      const newEnv = {
+        ...window.process.env,
+        LANGFLOW_AUTO_LOGIN: "false",
+        LANGFLOW_NEW_USER_IS_ACTIVE: "true",
+      };
+      Object.defineProperty(window.process, "env", {
+        value: newEnv,
+        writable: true,
+        configurable: true,
+      });
+      sessionStorage.setItem("testMockAutoLogin", "true");
+    });
+
+    // Create random usernames, passwords and flow names for the test
+    const userAName = "user_a_" + Math.random().toString(36).substring(5);
+    const userAPassword = "pass_a_" + Math.random().toString(36).substring(5);
+    const userAFlowName = "flow_a_" + Math.random().toString(36).substring(5);
+
+    // Log in as admin and create test user
+    await page.goto("/");
+    await page.waitForSelector("text=sign in to langflow", { timeout: 30000 });
+    await page.getByPlaceholder("Username").fill("langflow");
+    await page.getByPlaceholder("Password").fill("langflow");
+    await page.evaluate(() => {
+      sessionStorage.removeItem("testMockAutoLogin");
+    });
+    await page.getByRole("button", { name: "Sign In" }).click();
+
+    // Create User A
+    await page.waitForSelector('[data-testid="mainpage_title"]', {
+      timeout: 30000,
+    });
+    await page.getByTestId("user-profile-settings").click();
+    await page.getByText("Admin Page", { exact: true }).click();
+    await page.getByText("New User", { exact: true }).click();
+    await page.getByPlaceholder("Username").last().fill(userAName);
+    await page.locator('input[name="password"]').fill(userAPassword);
+    await page.locator('input[name="confirmpassword"]').fill(userAPassword);
+    await page.waitForSelector("#is_active", { timeout: 1500 });
+    await page.locator("#is_active").click();
+    await expect(page.locator("#is_active")).toBeChecked();
+    await page.getByText("Save", { exact: true }).click();
+    await page.waitForSelector("text=new user added", { timeout: 30000 });
+
+    // Log out from admin
+    await page.getByTestId("icon-ChevronLeft").first().click();
+    await page.waitForSelector("[data-testid='user-profile-settings']", {
+      timeout: 1500,
+    });
+    await page.getByTestId("user-profile-settings").click();
+    await page.evaluate(() => {
+      sessionStorage.setItem("testMockAutoLogin", "true");
+    });
+    await page.getByText("Logout", { exact: true }).click();
+
+    // ---- USER A SESSION ----
+
+    // Log in as User A
+    await page.waitForSelector("text=sign in to langflow", { timeout: 30000 });
+    await page.getByPlaceholder("Username").fill(userAName);
+    await page.getByPlaceholder("Password").fill(userAPassword);
+    await page.evaluate(() => {
+      sessionStorage.removeItem("testMockAutoLogin");
+    });
+    await page.getByRole("button", { name: "Sign In" }).click();
+
+    // Create a flow for User A
+    await page.waitForSelector('[id="new-project-btn"]', { timeout: 30000 });
+    // Check that User A starts with an empty flows list
+    expect(
+      (
+        await page.waitForSelector(
+          "text=Begin with a template, or start from scratch.",
+          { timeout: 30000 },
+        )
+      ).isVisible(),
+    );
+
+    await page.waitForSelector('[data-testid="mainpage_title"]', {
+      timeout: 30000,
+    });
+
+    try {
+      await page.getByTestId("new_project_btn_empty_page").click();
+    } catch (error) {
+      await page.getByText("New Flow", { exact: true }).click();
+    }
+
+    await page.waitForSelector('[data-testid="modal-title"]', {
+      timeout: 3000,
+    });
+    await page.getByRole("heading", { name: "Basic Prompting" }).click();
+    await page.waitForSelector('[data-testid="fit_view"]', { timeout: 30000 });
+    await page.getByTestId("flow_menu_trigger").click();
+    await page.getByText("Edit Details", { exact: true }).last().click();
+    await page.getByPlaceholder("Flow Name").fill(userAFlowName);
+    await page.getByText("Save", { exact: true }).click();
+    await page.waitForSelector('[data-testid="icon-ChevronLeft"]', {
+      timeout: 30000,
+    });
+    await page.getByTestId("icon-ChevronLeft").first().click();
+
+    // Verify User A can see their flow
+    await page.waitForSelector('[data-testid="search-store-input"]:enabled', {
+      timeout: 30000,
+    });
+    await expect(page.getByText(userAFlowName, { exact: true })).toBeVisible({
+      timeout: 2000,
+    });
+
+    // Log out User A
+    await page.getByTestId("user-profile-settings").click();
+    await page.evaluate(() => {
+      sessionStorage.setItem("testMockAutoLogin", "true");
+    });
+    await page.getByText("Logout", { exact: true }).click();
+
+    // ---- ADMIN SESSION AGAIN ----
+
+    // Log in as admin again
+    await page.waitForSelector("text=sign in to langflow", { timeout: 30000 });
+    await page.getByPlaceholder("Username").fill("langflow");
+    await page.getByPlaceholder("Password").fill("langflow");
+    await page.evaluate(() => {
+      sessionStorage.removeItem("testMockAutoLogin");
+    });
+    await page.getByRole("button", { name: "Sign In" }).click();
+
+    // Verify admin can't see User A's flow
+    await expect(page.getByText(userAFlowName, { exact: true })).toBeVisible({
+      timeout: 2000,
+      visible: false,
+    });
+
+    // Cleanup
+    await page.evaluate(() => {
+      sessionStorage.removeItem("testMockAutoLogin");
+    });
+  },
+);