Store user settings/data on the server and multi user support (#2160)

* wip per user data

* Rename, hide menu

* better error
rework default user

* store pretty

* Add userdata endpoints
Change nodetemplates to userdata

* add multi user message

* make normal arg

* Fix tests

* Ignore user dir

* user tests

* Changed to default to browser storage and add server-storage arg

* fix crash on empty templates

* fix settings added before load

* ignore parse errors

tests-ui/babel.config.json

@@ -1,3 +1,4 @@
 {
-	"presets": ["@babel/preset-env"]
+	"presets": ["@babel/preset-env"],
+	"plugins": ["babel-plugin-transform-import-meta"]
 }

tests-ui/package-lock.json

@@ -11,6 +11,7 @@
       "devDependencies": {
         "@babel/preset-env": "^7.22.20",
         "@types/jest": "^29.5.5",
+        "babel-plugin-transform-import-meta": "^2.2.1",
         "jest": "^29.7.0",
         "jest-environment-jsdom": "^29.7.0"
       }
@@ -2591,6 +2592,19 @@
         "@babel/core": "^7.4.0 || ^8.0.0-0 <8.0.0"
       }
     },
+    "node_modules/babel-plugin-transform-import-meta": {
+      "version": "2.2.1",
+      "resolved": "https://registry.npmjs.org/babel-plugin-transform-import-meta/-/babel-plugin-transform-import-meta-2.2.1.tgz",
+      "integrity": "sha512-AxNh27Pcg8Kt112RGa3Vod2QS2YXKKJ6+nSvRtv7qQTJAdx0MZa4UHZ4lnxHUWA2MNbLuZQv5FVab4P1CoLOWw==",
+      "dev": true,
+      "dependencies": {
+        "@babel/template": "^7.4.4",
+        "tslib": "^2.4.0"
+      },
+      "peerDependencies": {
+        "@babel/core": "^7.10.0"
+      }
+    },
     "node_modules/babel-preset-current-node-syntax": {
       "version": "1.0.1",
       "resolved": "https://registry.npmjs.org/babel-preset-current-node-syntax/-/babel-preset-current-node-syntax-1.0.1.tgz",
@@ -5233,6 +5247,12 @@
         "node": ">=12"
       }
     },
+    "node_modules/tslib": {
+      "version": "2.6.2",
+      "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.6.2.tgz",
+      "integrity": "sha512-AEYxH93jGFPn/a2iVAwW87VuUIkR1FVUKB77NwMF7nBTDkDrrT/Hpt/IrCJ0QXhW27jTBDcf5ZY7w6RiqTMw2Q==",
+      "dev": true
+    },
     "node_modules/type-detect": {
       "version": "4.0.8",
       "resolved": "https://registry.npmjs.org/type-detect/-/type-detect-4.0.8.tgz",

tests-ui/package.json

@@ -24,6 +24,7 @@
   "devDependencies": {
     "@babel/preset-env": "^7.22.20",
     "@types/jest": "^29.5.5",
+    "babel-plugin-transform-import-meta": "^2.2.1",
     "jest": "^29.7.0",
     "jest-environment-jsdom": "^29.7.0"
   }

tests-ui/tests/users.test.js

@@ -0,0 +1,295 @@
+// @ts-check
+/// <reference path="../node_modules/@types/jest/index.d.ts" />
+const { start } = require("../utils");
+const lg = require("../utils/litegraph");
+
+describe("users", () => {
+	beforeEach(() => {
+		lg.setup(global);
+	});
+
+	afterEach(() => {
+		lg.teardown(global);
+	});
+
+	function expectNoUserScreen() {
+		// Ensure login isnt visible
+		const selection = document.querySelectorAll("#comfy-user-selection")?.[0];
+		expect(selection["style"].display).toBe("none");
+		const menu = document.querySelectorAll(".comfy-menu")?.[0];
+		expect(window.getComputedStyle(menu)?.display).not.toBe("none");
+	}
+
+	describe("multi-user", () => {
+		function mockAddStylesheet() {
+			const utils = require("../../web/scripts/utils");
+			utils.addStylesheet = jest.fn().mockReturnValue(Promise.resolve());
+		}
+
+		async function waitForUserScreenShow() {
+			mockAddStylesheet();
+
+			// Wait for "show" to be called
+			const { UserSelectionScreen } = require("../../web/scripts/ui/userSelection");
+			let resolve, reject;
+			const fn = UserSelectionScreen.prototype.show;
+			const p = new Promise((res, rej) => {
+				resolve = res;
+				reject = rej;
+			});
+			jest.spyOn(UserSelectionScreen.prototype, "show").mockImplementation(async (...args) => {
+				const res = fn(...args);
+				await new Promise(process.nextTick); // wait for promises to resolve
+				resolve();
+				return res;
+			});
+			// @ts-ignore
+			setTimeout(() => reject("timeout waiting for UserSelectionScreen to be shown."), 500);
+			await p;
+			await new Promise(process.nextTick); // wait for promises to resolve
+		}
+
+		async function testUserScreen(onShown, users) {
+			if (!users) {
+				users = {};
+			}
+			const starting = start({
+				resetEnv: true,
+				userConfig: { storage: "server", users },
+			});
+
+			// Ensure no current user
+			expect(localStorage["Comfy.userId"]).toBeFalsy();
+			expect(localStorage["Comfy.userName"]).toBeFalsy();
+
+			await waitForUserScreenShow();
+
+			const selection = document.querySelectorAll("#comfy-user-selection")?.[0];
+			expect(selection).toBeTruthy();
+
+			// Ensure login is visible
+			expect(window.getComputedStyle(selection)?.display).not.toBe("none");
+			// Ensure menu is hidden
+			const menu = document.querySelectorAll(".comfy-menu")?.[0];
+			expect(window.getComputedStyle(menu)?.display).toBe("none");
+
+			const isCreate = await onShown(selection);
+
+			// Submit form
+			selection.querySelectorAll("form")[0].submit();
+			await new Promise(process.nextTick); // wait for promises to resolve
+
+			// Wait for start
+			const s = await starting;
+
+			// Ensure login is removed
+			expect(document.querySelectorAll("#comfy-user-selection")).toHaveLength(0);
+			expect(window.getComputedStyle(menu)?.display).not.toBe("none");
+
+			// Ensure settings + templates are saved
+			const { api } = require("../../web/scripts/api");
+			expect(api.createUser).toHaveBeenCalledTimes(+isCreate);
+			expect(api.storeSettings).toHaveBeenCalledTimes(+isCreate);
+			expect(api.storeUserData).toHaveBeenCalledTimes(+isCreate);
+			if (isCreate) {
+				expect(api.storeUserData).toHaveBeenCalledWith("comfy.templates.json", null, { stringify: false });
+				expect(s.app.isNewUserSession).toBeTruthy();
+			} else {
+				expect(s.app.isNewUserSession).toBeFalsy();
+			}
+
+			return { users, selection, ...s };
+		}
+
+		it("allows user creation if no users", async () => {
+			const { users } = await testUserScreen((selection) => {
+				// Ensure we have no users flag added
+				expect(selection.classList.contains("no-users")).toBeTruthy();
+
+				// Enter a username
+				const input = selection.getElementsByTagName("input")[0];
+				input.focus();
+				input.value = "Test User";
+
+				return true;
+			});
+
+			expect(users).toStrictEqual({
+				"Test User!": "Test User",
+			});
+
+			expect(localStorage["Comfy.userId"]).toBe("Test User!");
+			expect(localStorage["Comfy.userName"]).toBe("Test User");
+		});
+		it("allows user creation if no current user but other users", async () => {
+			const users = {
+				"Test User 2!": "Test User 2",
+			};
+
+			await testUserScreen((selection) => {
+				expect(selection.classList.contains("no-users")).toBeFalsy();
+
+				// Enter a username
+				const input = selection.getElementsByTagName("input")[0];
+				input.focus();
+				input.value = "Test User 3";
+				return true;
+			}, users);
+
+			expect(users).toStrictEqual({
+				"Test User 2!": "Test User 2",
+				"Test User 3!": "Test User 3",
+			});
+
+			expect(localStorage["Comfy.userId"]).toBe("Test User 3!");
+			expect(localStorage["Comfy.userName"]).toBe("Test User 3");
+		});
+		it("allows user selection if no current user but other users", async () => {
+			const users = {
+				"A!": "A",
+				"B!": "B",
+				"C!": "C",
+			};
+
+			await testUserScreen((selection) => {
+				expect(selection.classList.contains("no-users")).toBeFalsy();
+
+				// Check user list
+				const select = selection.getElementsByTagName("select")[0];
+				const options = select.getElementsByTagName("option");
+				expect(
+					[...options]
+						.filter((o) => !o.disabled)
+						.reduce((p, n) => {
+							p[n.getAttribute("value")] = n.textContent;
+							return p;
+						}, {})
+				).toStrictEqual(users);
+
+				// Select an option
+				select.focus();
+				select.value = options[2].value;
+
+				return false;
+			}, users);
+
+			expect(users).toStrictEqual(users);
+
+			expect(localStorage["Comfy.userId"]).toBe("B!");
+			expect(localStorage["Comfy.userName"]).toBe("B");
+		});
+		it("doesnt show user screen if current user", async () => {
+			const starting = start({
+				resetEnv: true,
+				userConfig: {
+					storage: "server",
+					users: {
+						"User!": "User",
+					},
+				},
+				localStorage: {
+					"Comfy.userId": "User!",
+					"Comfy.userName": "User",
+				},
+			});
+			await new Promise(process.nextTick); // wait for promises to resolve
+
+			expectNoUserScreen();
+
+			await starting;
+		});
+		it("allows user switching", async () => {
+			const { app } = await start({
+				resetEnv: true,
+				userConfig: {
+					storage: "server",
+					users: {
+						"User!": "User",
+					},
+				},
+				localStorage: {
+					"Comfy.userId": "User!",
+					"Comfy.userName": "User",
+				},
+			});
+
+			// cant actually test switching user easily but can check the setting is present
+			expect(app.ui.settings.settingsLookup["Comfy.SwitchUser"]).toBeTruthy();
+		});
+	});
+	describe("single-user", () => {
+		it("doesnt show user creation if no default user", async () => {
+			const { app } = await start({
+				resetEnv: true,
+				userConfig: { migrated: false, storage: "server" },
+			});
+			expectNoUserScreen();
+
+			// It should store the settings
+			const { api } = require("../../web/scripts/api");
+			expect(api.storeSettings).toHaveBeenCalledTimes(1);
+			expect(api.storeUserData).toHaveBeenCalledTimes(1);
+			expect(api.storeUserData).toHaveBeenCalledWith("comfy.templates.json", null, { stringify: false });
+			expect(app.isNewUserSession).toBeTruthy();
+		});
+		it("doesnt show user creation if default user", async () => {
+			const { app } = await start({
+				resetEnv: true,
+				userConfig: { migrated: true, storage: "server" },
+			});
+			expectNoUserScreen();
+
+			// It should store the settings
+			const { api } = require("../../web/scripts/api");
+			expect(api.storeSettings).toHaveBeenCalledTimes(0);
+			expect(api.storeUserData).toHaveBeenCalledTimes(0);
+			expect(app.isNewUserSession).toBeFalsy();
+		});
+		it("doesnt allow user switching", async () => {
+			const { app } = await start({
+				resetEnv: true,
+				userConfig: { migrated: true, storage: "server" },
+			});
+			expectNoUserScreen();
+
+			expect(app.ui.settings.settingsLookup["Comfy.SwitchUser"]).toBeFalsy();
+		});
+	});
+	describe("browser-user", () => {
+		it("doesnt show user creation if no default user", async () => {
+			const { app } = await start({
+				resetEnv: true,
+				userConfig: { migrated: false, storage: "browser" },
+			});
+			expectNoUserScreen();
+
+			// It should store the settings
+			const { api } = require("../../web/scripts/api");
+			expect(api.storeSettings).toHaveBeenCalledTimes(0);
+			expect(api.storeUserData).toHaveBeenCalledTimes(0);
+			expect(app.isNewUserSession).toBeFalsy();
+		});
+		it("doesnt show user creation if default user", async () => {
+			const { app } = await start({
+				resetEnv: true,
+				userConfig: { migrated: true, storage: "server" },
+			});
+			expectNoUserScreen();
+
+			// It should store the settings
+			const { api } = require("../../web/scripts/api");
+			expect(api.storeSettings).toHaveBeenCalledTimes(0);
+			expect(api.storeUserData).toHaveBeenCalledTimes(0);
+			expect(app.isNewUserSession).toBeFalsy();
+		});
+		it("doesnt allow user switching", async () => {
+			const { app } = await start({
+				resetEnv: true,
+				userConfig: { migrated: true, storage: "browser" },
+			});
+			expectNoUserScreen();
+
+			expect(app.ui.settings.settingsLookup["Comfy.SwitchUser"]).toBeFalsy();
+		});
+	});
+});

tests-ui/utils/index.js

@@ -1,10 +1,18 @@
 const { mockApi } = require("./setup");
 const { Ez } = require("./ezgraph");
 const lg = require("./litegraph");
+const fs = require("fs");
+const path = require("path");
+
+const html = fs.readFileSync(path.resolve(__dirname, "../../web/index.html"))
 
 /**
  *
- * @param { Parameters<mockApi>[0] & { resetEnv?: boolean, preSetup?(app): Promise<void> } } config
+ * @param { Parameters<typeof mockApi>[0] & { 
+ * 	resetEnv?: boolean, 
+ * 	preSetup?(app): Promise<void>,
+ *  localStorage?: Record<string, string> 
+ * } } config
  * @returns
  */
 export async function start(config = {}) {
@@ -12,12 +20,18 @@ export async function start(config = {}) {
 		jest.resetModules();
 		jest.resetAllMocks();
         lg.setup(global);
+		localStorage.clear();
+		sessionStorage.clear();
 	}
 
+	Object.assign(localStorage, config.localStorage ?? {});
+	document.body.innerHTML = html;
+
 	mockApi(config);
 	const { app } = require("../../web/scripts/app");
 	config.preSetup?.(app);
 	await app.setup();
+
 	return { ...Ez.graph(app, global["LiteGraph"], global["LGraphCanvas"]), app };
 }
 

tests-ui/utils/setup.js

@@ -18,9 +18,21 @@ function* walkSync(dir) {
  */
 
 /**
- * @param { { mockExtensions?: string[], mockNodeDefs?: Record<string, ComfyObjectInfo> } } config
+ * @param {{ 
+ * 	mockExtensions?: string[], 
+ * 	mockNodeDefs?: Record<string, ComfyObjectInfo>,
+* 	settings?: Record<string, string>
+* 	userConfig?: {storage: "server" | "browser", users?: Record<string, any>, migrated?: boolean },
+* 	userData?: Record<string, any>
+ * }} config
  */
-export function mockApi({ mockExtensions, mockNodeDefs } = {}) {
+export function mockApi(config = {}) {
+	let { mockExtensions, mockNodeDefs, userConfig, settings, userData } = {
+		userConfig,
+		settings: {},
+		userData: {},
+		...config,
+	};
 	if (!mockExtensions) {
 		mockExtensions = Array.from(walkSync(path.resolve("../web/extensions/core")))
 			.filter((x) => x.endsWith(".js"))
@@ -40,6 +52,26 @@ export function mockApi({ mockExtensions, mockNodeDefs } = {}) {
 		getNodeDefs: jest.fn(() => mockNodeDefs),
 		init: jest.fn(),
 		apiURL: jest.fn((x) => "../../web/" + x),
+		createUser: jest.fn((username) => {
+			if(username in userConfig.users) {
+				return { status: 400, json: () => "Duplicate" }
+			}
+			userConfig.users[username + "!"] = username;
+			return { status: 200, json: () => username + "!" }
+		}),
+		getUserConfig: jest.fn(() => userConfig ?? { storage: "browser", migrated: false }),
+		getSettings: jest.fn(() => settings),
+		storeSettings: jest.fn((v) => Object.assign(settings, v)),
+		getUserData: jest.fn((f) => {
+			if (f in userData) {
+				return { status: 200, json: () => userData[f] };
+			} else {
+				return { status: 404 };
+			}
+		}),
+		storeUserData: jest.fn((file, data) => {
+			userData[file] = data;
+		}),
 	};
 	jest.mock("../../web/scripts/api", () => ({
 		get api() {