Rework user selection (#1728)

* Move to new route

* Convert to tailwind

* Basic style

* Add userStore

* nit

* nit

* nit

* Remove app.#setUser

* Route to user-select view

* Mock login

* Use primevue UI components

* handle create new user

* Remove legacy user selection

* Add logout button on side toolbar

* Add username to logout button tooltip

* Add playwright tests

* hide logout button in single user server

* nit

src/components/sidebar/SideToolbar.vue

@@ -12,6 +12,7 @@
         @click="onTabClick(tab)"
       />
       <div class="side-tool-bar-end">
+        <SidebarLogoutIcon v-if="userStore.isMultiUserServer" />
         <SidebarThemeToggleIcon />
         <SidebarSettingsToggleIcon />
       </div>
@@ -29,15 +30,18 @@
 import SidebarIcon from './SidebarIcon.vue'
 import SidebarThemeToggleIcon from './SidebarThemeToggleIcon.vue'
 import SidebarSettingsToggleIcon from './SidebarSettingsToggleIcon.vue'
+import SidebarLogoutIcon from './SidebarLogoutIcon.vue'
 import ExtensionSlot from '@/components/common/ExtensionSlot.vue'
 import { computed } from 'vue'
 import { useWorkspaceStore } from '@/stores/workspaceStore'
 import { useSettingStore } from '@/stores/settingStore'
 import type { SidebarTabExtension } from '@/types/extensionTypes'
 import { useKeybindingStore } from '@/stores/keybindingStore'
+import { useUserStore } from '@/stores/userStore'
 
 const workspaceStore = useWorkspaceStore()
 const settingStore = useSettingStore()
+const userStore = useUserStore()
 
 const teleportTarget = computed(() =>
   settingStore.get('Comfy.Sidebar.Location') === 'left'

src/components/sidebar/SidebarLogoutIcon.vue

@@ -0,0 +1,21 @@
+<template>
+  <SidebarIcon icon="pi pi-sign-out" :tooltip="tooltip" @click="logout" />
+</template>
+
+<script setup lang="ts">
+import SidebarIcon from './SidebarIcon.vue'
+import { useUserStore } from '@/stores/userStore'
+import { computed } from 'vue'
+import { useI18n } from 'vue-i18n'
+
+const { t } = useI18n()
+const userStore = useUserStore()
+
+const tooltip = computed(
+  () => `${t('sideToolbar.logout')} (${userStore.currentUser?.username})`
+)
+const logout = () => {
+  userStore.logout()
+  window.location.reload()
+}
+</script>

src/locales/en.ts

@@ -3,6 +3,13 @@ export default {
     title: 'Welcome to ComfyUI',
     getStarted: 'Get Started'
   },
+  userSelect: {
+    newUser: 'New user',
+    enterUsername: 'Enter a username',
+    existingUser: 'Existing user',
+    selectUser: 'Select a user',
+    next: 'Next'
+  },
   notSupported: {
     title: 'Your device is not supported',
     message: 'Only following devices are supported:',
@@ -147,6 +154,7 @@ export default {
   newFolder: 'New Folder',
   sideToolbar: {
     themeToggle: 'Toggle Theme',
+    logout: 'Logout',
     queue: 'Queue',
     nodeLibrary: 'Node Library',
     workflows: 'Workflows',

src/locales/ja.ts

@@ -148,6 +148,7 @@ export default {
   sideToolbar: {
     themeToggle: 'テーマの切り替え',
     queue: 'キュー',
+    logout: 'ログアウト',
     nodeLibrary: 'ノードライブラリ',
     workflows: 'ワークフロー',
     browseTemplates: 'サンプルテンプレートを表示',

src/locales/ru.ts

@@ -60,6 +60,7 @@ export default {
   sideToolbar: {
     themeToggle: 'Переключить тему',
     queue: 'Очередь',
+    logout: 'Выйти',
     nodeLibrary: 'Библиотека узлов',
     workflows: 'Рабочие процессы',
     browseTemplates: 'Просмотреть примеры шаблонов',

src/locales/zh.ts

@@ -61,6 +61,7 @@ export default {
   newFolder: '新建文件夹',
   sideToolbar: {
     themeToggle: '主题切换',
+    logout: '登出',
     queue: '队列',
     nodeLibrary: '节点库',
     workflows: '工作流',

src/router.ts

@@ -7,6 +7,7 @@ import {
 } from 'vue-router'
 import LayoutDefault from '@/views/layouts/LayoutDefault.vue'
 import { isElectron } from './utils/envUtil'
+import { useUserStore } from './stores/userStore'
 
 const isFileProtocol = window.location.protocol === 'file:'
 const basePath = isElectron() ? '/' : window.location.pathname
@@ -38,7 +39,21 @@ const router = createRouter({
         {
           path: '',
           name: 'GraphView',
-          component: () => import('@/views/GraphView.vue')
+          component: () => import('@/views/GraphView.vue'),
+          beforeEnter: async (to, from, next) => {
+            const userStore = useUserStore()
+            await userStore.initialize()
+            if (userStore.needsLogin) {
+              next('/user-select')
+            } else {
+              next()
+            }
+          }
+        },
+        {
+          path: 'user-select',
+          name: 'UserSelectView',
+          component: () => import('@/views/UserSelectView.vue')
         },
         {
           path: 'server-start',

src/scripts/api.ts

@@ -42,6 +42,9 @@ export class ComfyApi extends EventTarget {
    * The current client id from websocket status updates.
    */
   clientId?: string
+  /**
+   * The current user id.
+   */
   user: string
   socket: WebSocket | null = null
 
@@ -49,7 +52,6 @@ export class ComfyApi extends EventTarget {
 
   constructor() {
     super()
-    // api.user is set by ComfyApp.setup()
     this.user = ''
     this.api_host = location.host
     this.api_base = location.pathname.split('/').slice(0, -1).join('/')

src/scripts/app.ts

@@ -1711,67 +1711,6 @@ export class ComfyApp {
     )
   }
 
-  async #setUser() {
-    const userConfig = await api.getUserConfig()
-    // Return in single user mode.
-    if (userConfig.users === undefined) {
-      return
-    }
-
-    let user = localStorage['Comfy.userId']
-    const users = userConfig.users ?? {}
-    if (!user || !users[user]) {
-      // Lift spinner / BlockUI for user selection.
-      if (this.vueAppReady) useWorkspaceStore().spinner = false
-
-      // This will rarely be hit so move the loading to on demand
-      const { UserSelectionScreen } = await import('./ui/userSelection')
-
-      this.ui.menuContainer.style.display = 'none'
-      const { userId, username } = await new UserSelectionScreen().show(
-        users,
-        user
-      )
-      this.ui.menuContainer.style.display = ''
-
-      user = userId
-      localStorage['Comfy.userName'] = username
-      localStorage['Comfy.userId'] = user
-    }
-
-    api.user = user
-
-    this.ui.settings.addSetting({
-      id: 'Comfy.SwitchUser',
-      name: 'Switch User',
-      type: (name) => {
-        let currentUser = localStorage['Comfy.userName']
-        if (currentUser) {
-          currentUser = ` (${currentUser})`
-        }
-        return $el('tr', [
-          $el('td', [
-            $el('label', {
-              textContent: name
-            })
-          ]),
-          $el('td', [
-            $el('button', {
-              textContent: name + (currentUser ?? ''),
-              onclick: () => {
-                delete localStorage['Comfy.userId']
-                delete localStorage['Comfy.userName']
-                window.location.reload()
-              }
-            })
-          ])
-        ])
-      },
-      // TODO: Is that the correct default value?
-      defaultValue: undefined
-    })
-  }
-
   /**
    * Set up the app on the page
    */
@@ -1779,7 +1718,6 @@ export class ComfyApp {
     this.canvasEl = canvasEl
     // Show menu container for GraphView.
     this.ui.menuContainer.style.display = 'block'
-    await this.#setUser()
 
     this.resizeCanvas()
 

src/scripts/ui/userSelection.css

@@ -1,135 +0,0 @@
-.comfy-user-selection {
-    width: 100vw;
-    height: 100vh;
-    position: absolute;
-    top: 0;
-    left: 0;
-    z-index: 999;
-    display: flex;
-    align-items: center;
-    justify-content: center;
-    font-family: sans-serif;
-    background: linear-gradient(var(--tr-even-bg-color), var(--tr-odd-bg-color));
-}
-
-.comfy-user-selection-inner {
-    background: var(--comfy-menu-bg);
-    margin-top: -30vh;
-    padding: 20px 40px;
-    border-radius: 10px;
-    min-width: 365px;
-    position: relative;
-    box-shadow: 0 0 20px rgba(0, 0, 0, 0.3);
-}
-
-.comfy-user-selection-inner form {
-    width: 100%;
-    display: flex;
-    flex-direction: column;
-    align-items: center;
-}
-
-.comfy-user-selection-inner h1 {
-    margin: 10px 0 30px 0;
-    font-weight: normal;
-}
-
-.comfy-user-selection-inner label {
-    display: flex;
-    flex-direction: column;
-    width: 100%;
-}
-
-.comfy-user-selection input,
-.comfy-user-selection select {
-    background-color: var(--comfy-input-bg);
-    color: var(--input-text);
-    border: 0;
-    border-radius: 5px;
-    padding: 5px;
-    margin-top: 10px;
-}
-
-.comfy-user-selection input::placeholder {
-    color: var(--descrip-text);
-    opacity: 1;
-}
-
-.comfy-user-existing {
-    width: 100%;
-}
-
-.no-users .comfy-user-existing {
-    display: none;
-}
-
-.comfy-user-selection-inner .or-separator {
-    margin: 10px 0;
-    padding: 10px;
-    display: block;
-    text-align: center;
-    width: 100%;
-    color: var(--descrip-text);
-}
-
-.comfy-user-selection-inner .or-separator {
-    overflow: hidden;
-    text-align: center;
-    margin-left: -10px;
-}
-
-.comfy-user-selection-inner .or-separator::before,
-.comfy-user-selection-inner .or-separator::after {
-    content: "";
-    background-color: var(--border-color);
-    position: relative;
-    height: 1px;
-    vertical-align: middle;
-    display: inline-block;
-    width: calc(50% - 20px);
-    top: -1px;
-}
-
-.comfy-user-selection-inner .or-separator::before {
-    right: 10px;
-    margin-left: -50%;
-}
-
-.comfy-user-selection-inner .or-separator::after {
-    left: 10px;
-    margin-right: -50%;
-}
-
-.comfy-user-selection-inner section {
-    width: 100%;
-    padding: 10px;
-    margin: -10px;
-    transition: background-color 0.2s;
-}
-
-.comfy-user-selection-inner section.selected {
-    background: var(--border-color);
-    border-radius: 5px;
-}
-
-.comfy-user-selection-inner footer {
-    display: flex;
-    flex-direction: column;
-    align-items: center;
-    margin-top: 20px;
-}
-
-.comfy-user-selection-inner .comfy-user-error {
-    color: var(--error-text);
-    margin-bottom: 10px;
-}
-
-.comfy-user-button-next {
-    font-size: 16px;
-    padding: 6px 10px;
-    width: 100px;
-    display: flex;
-    gap: 5px;
-    align-items: center;
-    justify-content: center;
-}

src/scripts/ui/userSelection.ts

@@ -1,146 +0,0 @@
-// @ts-strict-ignore
-import { api } from '../api'
-import { $el } from '../ui'
-import { createSpinner } from './spinner'
-import './userSelection.css'
-
-interface SelectedUser {
-  username: string
-  userId: string
-  created: boolean
-}
-
-export class UserSelectionScreen {
-  async show(users, user): Promise<SelectedUser> {
-    const userSelection = document.getElementById('comfy-user-selection')
-    userSelection.style.display = ''
-    return new Promise((resolve) => {
-      const input = userSelection.getElementsByTagName('input')[0]
-      const select = userSelection.getElementsByTagName('select')[0]
-      const inputSection = input.closest('section')
-      const selectSection = select.closest('section')
-      const form = userSelection.getElementsByTagName('form')[0]
-      const error = userSelection.getElementsByClassName('comfy-user-error')[0]
-      const button = userSelection.getElementsByClassName(
-        'comfy-user-button-next'
-      )[0]
-
-      let inputActive = null
-      input.addEventListener('focus', () => {
-        inputSection.classList.add('selected')
-        selectSection.classList.remove('selected')
-        inputActive = true
-      })
-      select.addEventListener('focus', () => {
-        inputSection.classList.remove('selected')
-        selectSection.classList.add('selected')
-        inputActive = false
-        select.style.color = ''
-      })
-      select.addEventListener('blur', () => {
-        if (!select.value) {
-          select.style.color = 'var(--descrip-text)'
-        }
-      })
-
-      form.addEventListener('submit', async (e) => {
-        e.preventDefault()
-        if (inputActive == null) {
-          error.textContent =
-            'Please enter a username or select an existing user.'
-        } else if (inputActive) {
-          const username = input.value.trim()
-          if (!username) {
-            error.textContent = 'Please enter a username.'
-            return
-          }
-
-          // Create new user
-          // Property 'readonly' does not exist on type 'HTMLSelectElement'.ts(2339)
-          // Property 'readonly' does not exist on type 'HTMLInputElement'. Did you mean 'readOnly'?ts(2551)
-          input.disabled =
-            select.disabled =
-            // @ts-expect-error
-            input.readonly =
-            // @ts-expect-error
-            select.readonly =
-              true
-          const spinner = createSpinner()
-          button.prepend(spinner)
-          try {
-            const resp = await api.createUser(username)
-            if (resp.status >= 300) {
-              let message =
-                'Error creating user: ' + resp.status + ' ' + resp.statusText
-              try {
-                const res = await resp.json()
-                if (res.error) {
-                  message = res.error
-                }
-              } catch (error) {}
-              throw new Error(message)
-            }
-
-            resolve({ username, userId: await resp.json(), created: true })
-          } catch (err) {
-            spinner.remove()
-            error.textContent =
-              err.message ??
-              err.statusText ??
-              err ??
-              'An unknown error occurred.'
-            // Property 'readonly' does not exist on type 'HTMLSelectElement'.ts(2339)
-            // Property 'readonly' does not exist on type 'HTMLInputElement'. Did you mean 'readOnly'?ts(2551)
-            input.disabled =
-              select.disabled =
-              // @ts-expect-error
-              input.readonly =
-              // @ts-expect-error
-              select.readonly =
-                false
-            return
-          }
-        } else if (!select.value) {
-          error.textContent = 'Please select an existing user.'
-          return
-        } else {
-          resolve({
-            username: users[select.value],
-            userId: select.value,
-            created: false
-          })
-        }
-      })
-
-      if (user) {
-        const name = localStorage['Comfy.userName']
-        if (name) {
-          input.value = name
-        }
-      }
-      if (input.value) {
-        // Focus the input, do this separately as sometimes browsers like to fill in the value
-        input.focus()
-      }
-
-      const userIds = Object.keys(users ?? {})
-      if (userIds.length) {
-        for (const u of userIds) {
-          $el('option', { textContent: users[u], value: u, parent: select })
-        }
-        select.style.color = 'var(--descrip-text)'
-
-        if (select.value) {
-          // Focus the select, do this separately as sometimes browsers like to fill in the value
-          select.focus()
-        }
-      } else {
-        userSelection.classList.add('no-users')
-        input.focus()
-      }
-    }).then((r: SelectedUser) => {
-      userSelection.remove()
-      return r
-    })
-  }
-}

src/stores/userStore.ts

@@ -0,0 +1,109 @@
+import { api } from '@/scripts/api'
+import { defineStore } from 'pinia'
+import { computed, ref, watchEffect } from 'vue'
+import type { User as UserConfig } from '@/types/apiTypes'
+
+export interface User {
+  userId: string
+  username: string
+}
+
+export const useUserStore = defineStore('user', () => {
+  /**
+   * The user config. null if not loaded.
+   */
+  const userConfig = ref<UserConfig | null>(null)
+  /**
+   * The current user id. null if not logged in or in single user mode.
+   */
+  const currentUserId = ref<string | null>(null)
+  const isMultiUserServer = computed(
+    () => userConfig.value && 'users' in userConfig.value
+  )
+  const needsLogin = computed(
+    () => !currentUserId.value && isMultiUserServer.value
+  )
+  const users = computed<User[]>(() =>
+    Object.entries(userConfig.value?.users ?? {}).map(([userId, username]) => ({
+      userId,
+      username
+    }))
+  )
+  const currentUser = computed<User | null>(
+    () =>
+      users.value.find((user) => user.userId === currentUserId.value) ?? null
+  )
+  const initialized = computed(() => userConfig.value !== null)
+
+  /**
+   * Initialize the user store.
+   */
+  async function initialize() {
+    userConfig.value = await api.getUserConfig()
+    currentUserId.value = localStorage['Comfy.userId']
+  }
+
+  /**
+   * Create a new user.
+   *
+   * @param username - The username.
+   * @returns The new user.
+   */
+  async function createUser(username: string): Promise<User> {
+    const resp = await api.createUser(username)
+    const data = await resp.json()
+    if (resp.status >= 300) {
+      throw new Error(
+        data.error ??
+          'Error creating user: ' + resp.status + ' ' + resp.statusText
+      )
+    }
+    return {
+      userId: data,
+      username
+    }
+  }
+
+  /**
+   * Login the current user.
+   *
+   * @param user - The user.
+   */
+  async function login({
+    userId,
+    username
+  }: {
+    userId: string
+    username: string
+  }) {
+    currentUserId.value = userId
+    localStorage['Comfy.userId'] = userId
+    localStorage['Comfy.userName'] = username
+  }
+
+  watchEffect(() => {
+    if (isMultiUserServer.value && currentUserId.value) {
+      api.user = currentUserId.value
+    }
+  })
+
+  /**
+   * Logout the current user.
+   */
+  async function logout() {
+    delete localStorage['Comfy.userId']
+    delete localStorage['Comfy.userName']
+  }
+
+  return {
+    users,
+    currentUser,
+    isMultiUserServer,
+    needsLogin,
+    initialized,
+    initialize,
+    createUser,
+    login,
+    logout
+  }
+})

src/types/apiTypes.ts

@@ -431,8 +431,10 @@ export const zSystemStats = z.object({
 })
 const zUser = z.object({
   storage: z.enum(['server']),
-  migrated: z.boolean(),
-  users: z.record(z.string(), z.unknown())
+  // `migrated` is only available in single-user mode.
+  migrated: z.boolean().optional(),
+  // `users` is only available in multi-user server mode.
+  users: z.record(z.string(), z.string()).optional()
 })
 const zUserData = z.array(z.array(z.string(), z.string()))
 const zUserDataFullInfo = z.object({

src/views/UserSelectView.vue

@@ -0,0 +1,90 @@
+<template>
+  <div
+    id="comfy-user-selection"
+    class="font-sans flex flex-col items-center h-screen m-0 text-neutral-300 bg-neutral-900 dark-theme pointer-events-auto"
+  >
+    <main
+      class="mt-[5vh] 2xl:mt-[20vh] min-w-84 relative rounded-lg bg-[var(--comfy-menu-bg)] p-5 px-10 shadow-lg"
+    >
+      <h1 class="my-2.5 mb-7 font-normal">ComfyUI</h1>
+      <form class="flex w-full flex-col items-center">
+        <div class="flex w-full flex-col gap-2">
+          <label for="new-user-input">{{ $t('userSelect.newUser') }}:</label>
+          <InputText
+            id="new-user-input"
+            v-model="newUsername"
+            :placeholder="$t('userSelect.enterUsername')"
+          />
+        </div>
+        <Divider />
+        <div class="flex w-full flex-col gap-2">
+          <label for="existing-user-select"
+            >{{ $t('userSelect.existingUser') }}:</label
+          >
+          <Select
+            v-model="selectedUser"
+            class="w-full"
+            inputId="existing-user-select"
+            :options="userStore.users"
+            option-label="username"
+            :placeholder="$t('userSelect.selectUser')"
+            :disabled="createNewUser"
+          />
+          <Message v-if="error" severity="error">{{ error }}</Message>
+        </div>
+        <footer class="mt-5">
+          <Button :label="$t('userSelect.next')" @click="login" />
+        </footer>
+      </form>
+    </main>
+  </div>
+</template>
+
+<script setup lang="ts">
+import Button from 'primevue/button'
+import Divider from 'primevue/divider'
+import InputText from 'primevue/inputtext'
+import Select from 'primevue/select'
+import Message from 'primevue/message'
+import { User, useUserStore } from '@/stores/userStore'
+import { useRouter } from 'vue-router'
+import { computed, onMounted, ref } from 'vue'
+
+const userStore = useUserStore()
+const router = useRouter()
+
+const selectedUser = ref<User | null>(null)
+const newUsername = ref('')
+const loginError = ref('')
+
+const createNewUser = computed(() => newUsername.value.trim() !== '')
+const newUserExistsError = computed(() => {
+  return userStore.users.find((user) => user.username === newUsername.value)
+    ? `User "${newUsername.value}" already exists`
+    : ''
+})
+const error = computed(() => newUserExistsError.value || loginError.value)
+
+const login = async () => {
+  try {
+    const user = createNewUser.value
+      ? await userStore.createUser(newUsername.value)
+      : selectedUser.value
+
+    if (!user) {
+      throw new Error('No user selected')
+    }
+
+    userStore.login(user)
+    router.push('/')
+  } catch (err) {
+    loginError.value = err.message ?? JSON.stringify(err)
+  }
+}
+
+onMounted(async () => {
+  if (!userStore.initialized) {
+    await userStore.initialize()
+  }
+})
+</script>