diff --git a/src/composables/auth/useCurrentUser.ts b/src/composables/auth/useCurrentUser.ts index cc5634cae..b43191fbb 100644 --- a/src/composables/auth/useCurrentUser.ts +++ b/src/composables/auth/useCurrentUser.ts @@ -41,8 +41,8 @@ export const useCurrentUser = () => { whenever(() => authStore.tokenRefreshTrigger, callback) const onUserLogout = (callback: () => void) => { - watch(resolvedUserInfo, (user) => { - if (!user) callback() + watch(resolvedUserInfo, (user, prevUser) => { + if (prevUser && !user) callback() }) } diff --git a/src/composables/useFeatureFlags.ts b/src/composables/useFeatureFlags.ts index 4e7368bc6..136b7ccd1 100644 --- a/src/composables/useFeatureFlags.ts +++ b/src/composables/useFeatureFlags.ts @@ -17,7 +17,8 @@ export enum ServerFeatureFlag { ONBOARDING_SURVEY_ENABLED = 'onboarding_survey_enabled', HUGGINGFACE_MODEL_IMPORT_ENABLED = 'huggingface_model_import_enabled', LINEAR_TOGGLE_ENABLED = 'linear_toggle_enabled', - ASYNC_MODEL_UPLOAD_ENABLED = 'async_model_upload_enabled' + ASYNC_MODEL_UPLOAD_ENABLED = 'async_model_upload_enabled', + TEAM_WORKSPACES_ENABLED = 'team_workspaces_enabled' } /** @@ -92,6 +93,12 @@ export function useFeatureFlags() { false ) ) + }, + get teamWorkspacesEnabled() { + return ( + remoteConfig.value.team_workspaces_enabled ?? + api.getServerFeature(ServerFeatureFlag.TEAM_WORKSPACES_ENABLED, false) + ) } }) diff --git a/src/locales/en/main.json b/src/locales/en/main.json index 9fd1284fd..65189f5b8 100644 --- a/src/locales/en/main.json +++ b/src/locales/en/main.json @@ -2592,5 +2592,20 @@ "completed": "Completed", "failed": "Failed" } + }, + "workspace": { + "unsavedChanges": { + "title": "Unsaved Changes", + "message": "You have unsaved changes. Do you want to discard them and switch workspaces?" + } + }, + "workspaceAuth": { + "errors": { + "notAuthenticated": "You must be logged in to access workspaces", + "invalidFirebaseToken": "Authentication failed. Please try logging in again.", + "accessDenied": "You do not have access to this workspace", + "workspaceNotFound": "Workspace not found", + "tokenExchangeFailed": "Failed to authenticate with workspace: {error}" + } } } \ No newline at end of file diff --git a/src/platform/auth/session/useSessionCookie.ts b/src/platform/auth/session/useSessionCookie.ts index 49f6fec46..a919b59eb 100644 --- a/src/platform/auth/session/useSessionCookie.ts +++ b/src/platform/auth/session/useSessionCookie.ts @@ -1,5 +1,6 @@ -import { api } from '@/scripts/api' import { isCloud } from '@/platform/distribution/types' +import { remoteConfig } from '@/platform/remoteConfig/remoteConfig' +import { api } from '@/scripts/api' import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore' /** @@ -10,31 +11,59 @@ export const useSessionCookie = () => { /** * Creates or refreshes the session cookie. * Called after login and on token refresh. + * + * When team_workspaces_enabled is true, uses Firebase token directly + * (since getAuthHeader() returns workspace token which shouldn't be used for session creation). + * When disabled, uses getAuthHeader() for backward compatibility. */ const createSession = async (): Promise => { if (!isCloud) return - const authStore = useFirebaseAuthStore() - const authHeader = await authStore.getAuthHeader() + try { + const authStore = useFirebaseAuthStore() - if (!authHeader) { - throw new Error('No auth header available for session creation') - } + let authHeader: Record - const response = await fetch(api.apiURL('/auth/session'), { - method: 'POST', - credentials: 'include', - headers: { - ...authHeader, - 'Content-Type': 'application/json' + if (remoteConfig.value.team_workspaces_enabled) { + const firebaseToken = await authStore.getIdToken() + if (!firebaseToken) { + console.warn( + 'Failed to create session cookie:', + 'No Firebase token available for session creation' + ) + return + } + authHeader = { Authorization: `Bearer ${firebaseToken}` } + } else { + const header = await authStore.getAuthHeader() + if (!header) { + console.warn( + 'Failed to create session cookie:', + 'No auth header available for session creation' + ) + return + } + authHeader = header } - }) - if (!response.ok) { - const errorData = await response.json().catch(() => ({})) - throw new Error( - `Failed to create session: ${errorData.message || response.statusText}` - ) + const response = await fetch(api.apiURL('/auth/session'), { + method: 'POST', + credentials: 'include', + headers: { + ...authHeader, + 'Content-Type': 'application/json' + } + }) + + if (!response.ok) { + const errorData = await response.json().catch(() => ({})) + console.warn( + 'Failed to create session cookie:', + errorData.message || response.statusText + ) + } + } catch (error) { + console.warn('Failed to create session cookie:', error) } } @@ -45,16 +74,21 @@ export const useSessionCookie = () => { const deleteSession = async (): Promise => { if (!isCloud) return - const response = await fetch(api.apiURL('/auth/session'), { - method: 'DELETE', - credentials: 'include' - }) + try { + const response = await fetch(api.apiURL('/auth/session'), { + method: 'DELETE', + credentials: 'include' + }) - if (!response.ok) { - const errorData = await response.json().catch(() => ({})) - throw new Error( - `Failed to delete session: ${errorData.message || response.statusText}` - ) + if (!response.ok) { + const errorData = await response.json().catch(() => ({})) + console.warn( + 'Failed to delete session cookie:', + errorData.message || response.statusText + ) + } + } catch (error) { + console.warn('Failed to delete session cookie:', error) } } diff --git a/src/platform/auth/workspace/useWorkspaceAuth.test.ts b/src/platform/auth/workspace/useWorkspaceAuth.test.ts new file mode 100644 index 000000000..68f5a198c --- /dev/null +++ b/src/platform/auth/workspace/useWorkspaceAuth.test.ts @@ -0,0 +1,670 @@ +import { createPinia, setActivePinia, storeToRefs } from 'pinia' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import { + useWorkspaceAuthStore, + WorkspaceAuthError +} from '@/stores/workspaceAuthStore' + +import { WORKSPACE_STORAGE_KEYS } from './workspaceConstants' + +const mockGetIdToken = vi.fn() + +vi.mock('@/stores/firebaseAuthStore', () => ({ + useFirebaseAuthStore: () => ({ + getIdToken: mockGetIdToken + }) +})) + +vi.mock('@/scripts/api', () => ({ + api: { + apiURL: (route: string) => `https://api.example.com/api${route}` + } +})) + +vi.mock('@/i18n', () => ({ + t: (key: string) => key +})) + +const mockRemoteConfig = vi.hoisted(() => ({ + value: { + team_workspaces_enabled: true + } +})) + +vi.mock('@/platform/remoteConfig/remoteConfig', () => ({ + remoteConfig: mockRemoteConfig +})) + +const mockWorkspace = { + id: 'workspace-123', + name: 'Test Workspace', + type: 'team' as const +} + +const mockWorkspaceWithRole = { + ...mockWorkspace, + role: 'owner' as const +} + +const mockTokenResponse = { + token: 'workspace-token-abc', + expires_at: new Date(Date.now() + 3600 * 1000).toISOString(), + workspace: mockWorkspace, + role: 'owner' as const, + permissions: ['owner:*'] +} + +describe('useWorkspaceAuthStore', () => { + beforeEach(() => { + setActivePinia(createPinia()) + vi.clearAllMocks() + vi.useFakeTimers() + sessionStorage.clear() + }) + + afterEach(() => { + vi.useRealTimers() + }) + + describe('initial state', () => { + it('has correct initial state values', () => { + const store = useWorkspaceAuthStore() + const { + currentWorkspace, + workspaceToken, + isAuthenticated, + isLoading, + error + } = storeToRefs(store) + + expect(currentWorkspace.value).toBeNull() + expect(workspaceToken.value).toBeNull() + expect(isAuthenticated.value).toBe(false) + expect(isLoading.value).toBe(false) + expect(error.value).toBeNull() + }) + }) + + describe('initializeFromSession', () => { + it('returns true and populates state when valid session data exists', () => { + const futureExpiry = Date.now() + 3600 * 1000 + sessionStorage.setItem( + WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE, + JSON.stringify(mockWorkspaceWithRole) + ) + sessionStorage.setItem(WORKSPACE_STORAGE_KEYS.TOKEN, 'valid-token') + sessionStorage.setItem( + WORKSPACE_STORAGE_KEYS.EXPIRES_AT, + futureExpiry.toString() + ) + + const store = useWorkspaceAuthStore() + const { currentWorkspace, workspaceToken } = storeToRefs(store) + + const result = store.initializeFromSession() + + expect(result).toBe(true) + expect(currentWorkspace.value).toEqual(mockWorkspaceWithRole) + expect(workspaceToken.value).toBe('valid-token') + }) + + it('returns false when sessionStorage is empty', () => { + const store = useWorkspaceAuthStore() + + const result = store.initializeFromSession() + + expect(result).toBe(false) + }) + + it('returns false and clears storage when token is expired', () => { + const pastExpiry = Date.now() - 1000 + sessionStorage.setItem( + WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE, + JSON.stringify(mockWorkspaceWithRole) + ) + sessionStorage.setItem(WORKSPACE_STORAGE_KEYS.TOKEN, 'expired-token') + sessionStorage.setItem( + WORKSPACE_STORAGE_KEYS.EXPIRES_AT, + pastExpiry.toString() + ) + + const store = useWorkspaceAuthStore() + + const result = store.initializeFromSession() + + expect(result).toBe(false) + expect( + sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE) + ).toBeNull() + expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBeNull() + expect( + sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT) + ).toBeNull() + }) + + it('returns false and clears storage when data is malformed', () => { + sessionStorage.setItem( + WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE, + 'invalid-json{' + ) + sessionStorage.setItem(WORKSPACE_STORAGE_KEYS.TOKEN, 'some-token') + sessionStorage.setItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT, 'not-a-number') + + const store = useWorkspaceAuthStore() + + const result = store.initializeFromSession() + + expect(result).toBe(false) + expect( + sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE) + ).toBeNull() + expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBeNull() + expect( + sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT) + ).toBeNull() + }) + + it('returns false when partial session data exists (missing token)', () => { + sessionStorage.setItem( + WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE, + JSON.stringify(mockWorkspaceWithRole) + ) + sessionStorage.setItem( + WORKSPACE_STORAGE_KEYS.EXPIRES_AT, + (Date.now() + 3600 * 1000).toString() + ) + + const store = useWorkspaceAuthStore() + + const result = store.initializeFromSession() + + expect(result).toBe(false) + }) + }) + + describe('switchWorkspace', () => { + it('successfully exchanges Firebase token for workspace token', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue({ + ok: true, + json: () => Promise.resolve(mockTokenResponse) + }) + ) + + const store = useWorkspaceAuthStore() + const { currentWorkspace, workspaceToken, isAuthenticated } = + storeToRefs(store) + + await store.switchWorkspace('workspace-123') + + expect(currentWorkspace.value).toEqual(mockWorkspaceWithRole) + expect(workspaceToken.value).toBe('workspace-token-abc') + expect(isAuthenticated.value).toBe(true) + }) + + it('stores workspace data in sessionStorage', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue({ + ok: true, + json: () => Promise.resolve(mockTokenResponse) + }) + ) + + const store = useWorkspaceAuthStore() + + await store.switchWorkspace('workspace-123') + + expect( + sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE) + ).toBe(JSON.stringify(mockWorkspaceWithRole)) + expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBe( + 'workspace-token-abc' + ) + expect( + sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT) + ).toBeTruthy() + }) + + it('sets isLoading to true during operation', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + let resolveResponse: (value: unknown) => void + const responsePromise = new Promise((resolve) => { + resolveResponse = resolve + }) + vi.stubGlobal('fetch', vi.fn().mockReturnValue(responsePromise)) + + const store = useWorkspaceAuthStore() + const { isLoading } = storeToRefs(store) + + const switchPromise = store.switchWorkspace('workspace-123') + expect(isLoading.value).toBe(true) + + resolveResponse!({ + ok: true, + json: () => Promise.resolve(mockTokenResponse) + }) + await switchPromise + + expect(isLoading.value).toBe(false) + }) + + it('throws WorkspaceAuthError with code NOT_AUTHENTICATED when Firebase token unavailable', async () => { + mockGetIdToken.mockResolvedValue(undefined) + + const store = useWorkspaceAuthStore() + const { error } = storeToRefs(store) + + await expect(store.switchWorkspace('workspace-123')).rejects.toThrow( + WorkspaceAuthError + ) + + expect(error.value).toBeInstanceOf(WorkspaceAuthError) + expect((error.value as WorkspaceAuthError).code).toBe('NOT_AUTHENTICATED') + }) + + it('throws WorkspaceAuthError with code ACCESS_DENIED on 403 response', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue({ + ok: false, + status: 403, + statusText: 'Forbidden', + json: () => Promise.resolve({ message: 'Access denied' }) + }) + ) + + const store = useWorkspaceAuthStore() + const { error } = storeToRefs(store) + + await expect(store.switchWorkspace('workspace-123')).rejects.toThrow( + WorkspaceAuthError + ) + + expect(error.value).toBeInstanceOf(WorkspaceAuthError) + expect((error.value as WorkspaceAuthError).code).toBe('ACCESS_DENIED') + }) + + it('throws WorkspaceAuthError with code WORKSPACE_NOT_FOUND on 404 response', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue({ + ok: false, + status: 404, + statusText: 'Not Found', + json: () => Promise.resolve({ message: 'Workspace not found' }) + }) + ) + + const store = useWorkspaceAuthStore() + const { error } = storeToRefs(store) + + await expect(store.switchWorkspace('workspace-123')).rejects.toThrow( + WorkspaceAuthError + ) + + expect(error.value).toBeInstanceOf(WorkspaceAuthError) + expect((error.value as WorkspaceAuthError).code).toBe( + 'WORKSPACE_NOT_FOUND' + ) + }) + + it('throws WorkspaceAuthError with code INVALID_FIREBASE_TOKEN on 401 response', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue({ + ok: false, + status: 401, + statusText: 'Unauthorized', + json: () => Promise.resolve({ message: 'Invalid token' }) + }) + ) + + const store = useWorkspaceAuthStore() + const { error } = storeToRefs(store) + + await expect(store.switchWorkspace('workspace-123')).rejects.toThrow( + WorkspaceAuthError + ) + + expect(error.value).toBeInstanceOf(WorkspaceAuthError) + expect((error.value as WorkspaceAuthError).code).toBe( + 'INVALID_FIREBASE_TOKEN' + ) + }) + + it('throws WorkspaceAuthError with code TOKEN_EXCHANGE_FAILED on other errors', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue({ + ok: false, + status: 500, + statusText: 'Internal Server Error', + json: () => Promise.resolve({ message: 'Server error' }) + }) + ) + + const store = useWorkspaceAuthStore() + const { error } = storeToRefs(store) + + await expect(store.switchWorkspace('workspace-123')).rejects.toThrow( + WorkspaceAuthError + ) + + expect(error.value).toBeInstanceOf(WorkspaceAuthError) + expect((error.value as WorkspaceAuthError).code).toBe( + 'TOKEN_EXCHANGE_FAILED' + ) + }) + + it('sends correct request to API', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: () => Promise.resolve(mockTokenResponse) + }) + vi.stubGlobal('fetch', mockFetch) + + const store = useWorkspaceAuthStore() + + await store.switchWorkspace('workspace-123') + + expect(mockFetch).toHaveBeenCalledWith( + 'https://api.example.com/api/auth/token', + { + method: 'POST', + headers: { + Authorization: 'Bearer firebase-token-xyz', + 'Content-Type': 'application/json' + }, + body: JSON.stringify({ workspace_id: 'workspace-123' }) + } + ) + }) + }) + + describe('clearWorkspaceContext', () => { + it('clears all state refs', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue({ + ok: true, + json: () => Promise.resolve(mockTokenResponse) + }) + ) + + const store = useWorkspaceAuthStore() + const { currentWorkspace, workspaceToken, error, isAuthenticated } = + storeToRefs(store) + + await store.switchWorkspace('workspace-123') + expect(isAuthenticated.value).toBe(true) + + store.clearWorkspaceContext() + + expect(currentWorkspace.value).toBeNull() + expect(workspaceToken.value).toBeNull() + expect(error.value).toBeNull() + expect(isAuthenticated.value).toBe(false) + }) + + it('clears sessionStorage', async () => { + sessionStorage.setItem( + WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE, + JSON.stringify(mockWorkspaceWithRole) + ) + sessionStorage.setItem(WORKSPACE_STORAGE_KEYS.TOKEN, 'some-token') + sessionStorage.setItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT, '12345') + + const store = useWorkspaceAuthStore() + + store.clearWorkspaceContext() + + expect( + sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE) + ).toBeNull() + expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBeNull() + expect( + sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT) + ).toBeNull() + }) + }) + + describe('getWorkspaceAuthHeader', () => { + it('returns null when no workspace token', () => { + const store = useWorkspaceAuthStore() + + const header = store.getWorkspaceAuthHeader() + + expect(header).toBeNull() + }) + + it('returns proper Authorization header when workspace token exists', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue({ + ok: true, + json: () => Promise.resolve(mockTokenResponse) + }) + ) + + const store = useWorkspaceAuthStore() + + await store.switchWorkspace('workspace-123') + const header = store.getWorkspaceAuthHeader() + + expect(header).toEqual({ + Authorization: 'Bearer workspace-token-abc' + }) + }) + }) + + describe('token refresh scheduling', () => { + it('schedules token refresh 5 minutes before expiry', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + const expiresInMs = 3600 * 1000 + const tokenResponseWithFutureExpiry = { + ...mockTokenResponse, + expires_at: new Date(Date.now() + expiresInMs).toISOString() + } + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: () => Promise.resolve(tokenResponseWithFutureExpiry) + }) + vi.stubGlobal('fetch', mockFetch) + + const store = useWorkspaceAuthStore() + + await store.switchWorkspace('workspace-123') + + expect(mockFetch).toHaveBeenCalledTimes(1) + + const refreshBufferMs = 5 * 60 * 1000 + const refreshDelay = expiresInMs - refreshBufferMs + + vi.advanceTimersByTime(refreshDelay - 1) + expect(mockFetch).toHaveBeenCalledTimes(1) + + await vi.advanceTimersByTimeAsync(1) + + expect(mockFetch).toHaveBeenCalledTimes(2) + }) + + it('clears context when refresh fails with ACCESS_DENIED', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + const expiresInMs = 3600 * 1000 + const tokenResponseWithFutureExpiry = { + ...mockTokenResponse, + expires_at: new Date(Date.now() + expiresInMs).toISOString() + } + const mockFetch = vi + .fn() + .mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve(tokenResponseWithFutureExpiry) + }) + .mockResolvedValueOnce({ + ok: false, + status: 403, + statusText: 'Forbidden', + json: () => Promise.resolve({ message: 'Access denied' }) + }) + vi.stubGlobal('fetch', mockFetch) + + const store = useWorkspaceAuthStore() + const { currentWorkspace, workspaceToken } = storeToRefs(store) + + await store.switchWorkspace('workspace-123') + expect(workspaceToken.value).toBe('workspace-token-abc') + + const refreshBufferMs = 5 * 60 * 1000 + const refreshDelay = expiresInMs - refreshBufferMs + + vi.advanceTimersByTime(refreshDelay) + await vi.waitFor(() => { + expect(currentWorkspace.value).toBeNull() + }) + + expect(workspaceToken.value).toBeNull() + }) + }) + + describe('refreshToken', () => { + it('does nothing when no current workspace', async () => { + const mockFetch = vi.fn() + vi.stubGlobal('fetch', mockFetch) + + const store = useWorkspaceAuthStore() + + await store.refreshToken() + + expect(mockFetch).not.toHaveBeenCalled() + }) + + it('refreshes token for current workspace', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: () => Promise.resolve(mockTokenResponse) + }) + vi.stubGlobal('fetch', mockFetch) + + const store = useWorkspaceAuthStore() + const { workspaceToken } = storeToRefs(store) + + await store.switchWorkspace('workspace-123') + expect(mockFetch).toHaveBeenCalledTimes(1) + + mockFetch.mockResolvedValue({ + ok: true, + json: () => + Promise.resolve({ + ...mockTokenResponse, + token: 'refreshed-token' + }) + }) + + await store.refreshToken() + expect(mockFetch).toHaveBeenCalledTimes(2) + expect(workspaceToken.value).toBe('refreshed-token') + }) + }) + + describe('isAuthenticated computed', () => { + it('returns true when both workspace and token are present', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue({ + ok: true, + json: () => Promise.resolve(mockTokenResponse) + }) + ) + + const store = useWorkspaceAuthStore() + const { isAuthenticated } = storeToRefs(store) + + await store.switchWorkspace('workspace-123') + + expect(isAuthenticated.value).toBe(true) + }) + + it('returns false when workspace is null', () => { + const store = useWorkspaceAuthStore() + const { isAuthenticated } = storeToRefs(store) + + expect(isAuthenticated.value).toBe(false) + }) + + it('returns false when currentWorkspace is set but workspaceToken is null', async () => { + mockGetIdToken.mockResolvedValue(null) + + const store = useWorkspaceAuthStore() + const { currentWorkspace, workspaceToken, isAuthenticated } = + storeToRefs(store) + + currentWorkspace.value = mockWorkspaceWithRole + workspaceToken.value = null + + expect(isAuthenticated.value).toBe(false) + }) + }) + + describe('feature flag disabled', () => { + beforeEach(() => { + mockRemoteConfig.value.team_workspaces_enabled = false + }) + + afterEach(() => { + mockRemoteConfig.value.team_workspaces_enabled = true + }) + + it('initializeFromSession returns false when flag disabled', () => { + const futureExpiry = Date.now() + 3600 * 1000 + sessionStorage.setItem( + WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE, + JSON.stringify(mockWorkspaceWithRole) + ) + sessionStorage.setItem(WORKSPACE_STORAGE_KEYS.TOKEN, 'valid-token') + sessionStorage.setItem( + WORKSPACE_STORAGE_KEYS.EXPIRES_AT, + futureExpiry.toString() + ) + + const store = useWorkspaceAuthStore() + const { currentWorkspace, workspaceToken } = storeToRefs(store) + + const result = store.initializeFromSession() + + expect(result).toBe(false) + expect(currentWorkspace.value).toBeNull() + expect(workspaceToken.value).toBeNull() + }) + + it('switchWorkspace is a no-op when flag disabled', async () => { + mockGetIdToken.mockResolvedValue('firebase-token-xyz') + const mockFetch = vi.fn() + vi.stubGlobal('fetch', mockFetch) + + const store = useWorkspaceAuthStore() + const { currentWorkspace, workspaceToken, isLoading } = storeToRefs(store) + + await store.switchWorkspace('workspace-123') + + expect(mockFetch).not.toHaveBeenCalled() + expect(currentWorkspace.value).toBeNull() + expect(workspaceToken.value).toBeNull() + expect(isLoading.value).toBe(false) + }) + }) +}) diff --git a/src/platform/auth/workspace/useWorkspaceSwitch.test.ts b/src/platform/auth/workspace/useWorkspaceSwitch.test.ts new file mode 100644 index 000000000..ae71937b2 --- /dev/null +++ b/src/platform/auth/workspace/useWorkspaceSwitch.test.ts @@ -0,0 +1,166 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import { useWorkspaceSwitch } from '@/platform/auth/workspace/useWorkspaceSwitch' +import type { WorkspaceWithRole } from '@/platform/auth/workspace/workspaceTypes' + +const mockSwitchWorkspace = vi.hoisted(() => vi.fn()) +const mockCurrentWorkspace = vi.hoisted(() => ({ + value: null as WorkspaceWithRole | null +})) + +vi.mock('@/stores/workspaceAuthStore', () => ({ + useWorkspaceAuthStore: () => ({ + switchWorkspace: mockSwitchWorkspace + }) +})) + +vi.mock('pinia', () => ({ + storeToRefs: () => ({ + currentWorkspace: mockCurrentWorkspace + }) +})) + +const mockModifiedWorkflows = vi.hoisted( + () => [] as Array<{ isModified: boolean }> +) + +vi.mock('@/platform/workflow/management/stores/workflowStore', () => ({ + useWorkflowStore: () => ({ + get modifiedWorkflows() { + return mockModifiedWorkflows + } + }) +})) + +const mockConfirm = vi.hoisted(() => vi.fn()) + +vi.mock('@/services/dialogService', () => ({ + useDialogService: () => ({ + confirm: mockConfirm + }) +})) + +vi.mock('vue-i18n', () => ({ + useI18n: () => ({ + t: (key: string) => key + }) +})) + +const mockReload = vi.fn() + +describe('useWorkspaceSwitch', () => { + beforeEach(() => { + vi.clearAllMocks() + mockCurrentWorkspace.value = { + id: 'workspace-1', + name: 'Test Workspace', + type: 'personal', + role: 'owner' + } + mockModifiedWorkflows.length = 0 + vi.stubGlobal('location', { reload: mockReload }) + }) + + afterEach(() => { + vi.unstubAllGlobals() + }) + + describe('hasUnsavedChanges', () => { + it('returns true when there are modified workflows', () => { + mockModifiedWorkflows.push({ isModified: true }) + const { hasUnsavedChanges } = useWorkspaceSwitch() + + expect(hasUnsavedChanges()).toBe(true) + }) + + it('returns true when multiple workflows are modified', () => { + mockModifiedWorkflows.push({ isModified: true }, { isModified: true }) + const { hasUnsavedChanges } = useWorkspaceSwitch() + + expect(hasUnsavedChanges()).toBe(true) + }) + + it('returns false when no workflows are modified', () => { + mockModifiedWorkflows.length = 0 + const { hasUnsavedChanges } = useWorkspaceSwitch() + + expect(hasUnsavedChanges()).toBe(false) + }) + }) + + describe('switchWithConfirmation', () => { + it('returns true immediately if switching to the same workspace', async () => { + const { switchWithConfirmation } = useWorkspaceSwitch() + + const result = await switchWithConfirmation('workspace-1') + + expect(result).toBe(true) + expect(mockSwitchWorkspace).not.toHaveBeenCalled() + expect(mockConfirm).not.toHaveBeenCalled() + }) + + it('switches directly without dialog when no unsaved changes', async () => { + mockModifiedWorkflows.length = 0 + mockSwitchWorkspace.mockResolvedValue(undefined) + const { switchWithConfirmation } = useWorkspaceSwitch() + + const result = await switchWithConfirmation('workspace-2') + + expect(result).toBe(true) + expect(mockConfirm).not.toHaveBeenCalled() + expect(mockSwitchWorkspace).toHaveBeenCalledWith('workspace-2') + expect(mockReload).toHaveBeenCalled() + }) + + it('shows confirmation dialog when there are unsaved changes', async () => { + mockModifiedWorkflows.push({ isModified: true }) + mockConfirm.mockResolvedValue(true) + mockSwitchWorkspace.mockResolvedValue(undefined) + const { switchWithConfirmation } = useWorkspaceSwitch() + + await switchWithConfirmation('workspace-2') + + expect(mockConfirm).toHaveBeenCalledWith({ + title: 'workspace.unsavedChanges.title', + message: 'workspace.unsavedChanges.message', + type: 'dirtyClose' + }) + }) + + it('returns false if user cancels the confirmation dialog', async () => { + mockModifiedWorkflows.push({ isModified: true }) + mockConfirm.mockResolvedValue(false) + const { switchWithConfirmation } = useWorkspaceSwitch() + + const result = await switchWithConfirmation('workspace-2') + + expect(result).toBe(false) + expect(mockSwitchWorkspace).not.toHaveBeenCalled() + expect(mockReload).not.toHaveBeenCalled() + }) + + it('calls switchWorkspace and reloads page after user confirms', async () => { + mockModifiedWorkflows.push({ isModified: true }) + mockConfirm.mockResolvedValue(true) + mockSwitchWorkspace.mockResolvedValue(undefined) + const { switchWithConfirmation } = useWorkspaceSwitch() + + const result = await switchWithConfirmation('workspace-2') + + expect(result).toBe(true) + expect(mockSwitchWorkspace).toHaveBeenCalledWith('workspace-2') + expect(mockReload).toHaveBeenCalled() + }) + + it('returns false if switchWorkspace throws an error', async () => { + mockModifiedWorkflows.length = 0 + mockSwitchWorkspace.mockRejectedValue(new Error('Switch failed')) + const { switchWithConfirmation } = useWorkspaceSwitch() + + const result = await switchWithConfirmation('workspace-2') + + expect(result).toBe(false) + expect(mockReload).not.toHaveBeenCalled() + }) + }) +}) diff --git a/src/platform/auth/workspace/useWorkspaceSwitch.ts b/src/platform/auth/workspace/useWorkspaceSwitch.ts new file mode 100644 index 000000000..6fb970d7f --- /dev/null +++ b/src/platform/auth/workspace/useWorkspaceSwitch.ts @@ -0,0 +1,49 @@ +import { storeToRefs } from 'pinia' +import { useI18n } from 'vue-i18n' + +import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore' +import { useDialogService } from '@/services/dialogService' +import { useWorkspaceAuthStore } from '@/stores/workspaceAuthStore' + +export function useWorkspaceSwitch() { + const { t } = useI18n() + const workspaceAuthStore = useWorkspaceAuthStore() + const { currentWorkspace } = storeToRefs(workspaceAuthStore) + const workflowStore = useWorkflowStore() + const dialogService = useDialogService() + + function hasUnsavedChanges(): boolean { + return workflowStore.modifiedWorkflows.length > 0 + } + + async function switchWithConfirmation(workspaceId: string): Promise { + if (currentWorkspace.value?.id === workspaceId) { + return true + } + + if (hasUnsavedChanges()) { + const confirmed = await dialogService.confirm({ + title: t('workspace.unsavedChanges.title'), + message: t('workspace.unsavedChanges.message'), + type: 'dirtyClose' + }) + + if (!confirmed) { + return false + } + } + + try { + await workspaceAuthStore.switchWorkspace(workspaceId) + window.location.reload() + return true + } catch { + return false + } + } + + return { + hasUnsavedChanges, + switchWithConfirmation + } +} diff --git a/src/platform/auth/workspace/workspaceConstants.ts b/src/platform/auth/workspace/workspaceConstants.ts new file mode 100644 index 000000000..cc28d1f47 --- /dev/null +++ b/src/platform/auth/workspace/workspaceConstants.ts @@ -0,0 +1,7 @@ +export const WORKSPACE_STORAGE_KEYS = { + CURRENT_WORKSPACE: 'Comfy.Workspace.Current', + TOKEN: 'Comfy.Workspace.Token', + EXPIRES_AT: 'Comfy.Workspace.ExpiresAt' +} as const + +export const TOKEN_REFRESH_BUFFER_MS = 5 * 60 * 1000 diff --git a/src/platform/auth/workspace/workspaceTypes.ts b/src/platform/auth/workspace/workspaceTypes.ts new file mode 100644 index 000000000..30774aef3 --- /dev/null +++ b/src/platform/auth/workspace/workspaceTypes.ts @@ -0,0 +1,6 @@ +export interface WorkspaceWithRole { + id: string + name: string + type: 'personal' | 'team' + role: 'owner' | 'member' +} diff --git a/src/platform/remoteConfig/types.ts b/src/platform/remoteConfig/types.ts index 7de4c2da2..7b8b1721c 100644 --- a/src/platform/remoteConfig/types.ts +++ b/src/platform/remoteConfig/types.ts @@ -42,4 +42,5 @@ export type RemoteConfig = { huggingface_model_import_enabled?: boolean linear_toggle_enabled?: boolean async_model_upload_enabled?: boolean + team_workspaces_enabled?: boolean } diff --git a/src/stores/firebaseAuthStore.ts b/src/stores/firebaseAuthStore.ts index 9f3039887..baa1840a8 100644 --- a/src/stores/firebaseAuthStore.ts +++ b/src/stores/firebaseAuthStore.ts @@ -23,7 +23,9 @@ import { useFirebaseAuth } from 'vuefire' import { getComfyApiBaseUrl } from '@/config/comfyApi' import { t } from '@/i18n' +import { WORKSPACE_STORAGE_KEYS } from '@/platform/auth/workspace/workspaceConstants' import { isCloud } from '@/platform/distribution/types' +import { remoteConfig } from '@/platform/remoteConfig/remoteConfig' import { useTelemetry } from '@/platform/telemetry' import { useDialogService } from '@/services/dialogService' import { useApiKeyAuthStore } from '@/stores/apiKeyAuthStore' @@ -107,6 +109,15 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => { isInitialized.value = true if (user === null) { lastTokenUserId.value = null + + // Clear workspace sessionStorage on logout to prevent stale tokens + try { + sessionStorage.removeItem(WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE) + sessionStorage.removeItem(WORKSPACE_STORAGE_KEYS.TOKEN) + sessionStorage.removeItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT) + } catch { + // Ignore sessionStorage errors (e.g., in private browsing mode) + } } // Reset balance when auth state changes @@ -152,16 +163,34 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => { /** * Retrieves the appropriate authentication header for API requests. * Checks for authentication in the following order: - * 1. Firebase authentication token (if user is logged in) - * 2. API key (if stored in the browser's credential manager) + * 1. Workspace token (if team_workspaces_enabled and user has active workspace context) + * 2. Firebase authentication token (if user is logged in) + * 3. API key (if stored in the browser's credential manager) * * @returns {Promise} - * - A LoggedInAuthHeader with Bearer token if Firebase authenticated + * - A LoggedInAuthHeader with Bearer token (workspace or Firebase) * - An ApiKeyAuthHeader with X-API-KEY if API key exists - * - null if neither authentication method is available + * - null if no authentication method is available */ const getAuthHeader = async (): Promise => { - // If available, set header with JWT used to identify the user to Firebase service + if (remoteConfig.value.team_workspaces_enabled) { + const workspaceToken = sessionStorage.getItem( + WORKSPACE_STORAGE_KEYS.TOKEN + ) + const expiresAt = sessionStorage.getItem( + WORKSPACE_STORAGE_KEYS.EXPIRES_AT + ) + + if (workspaceToken && expiresAt) { + const expiryTime = parseInt(expiresAt, 10) + if (Date.now() < expiryTime) { + return { + Authorization: `Bearer ${workspaceToken}` + } + } + } + } + const token = await getIdToken() if (token) { return { @@ -169,7 +198,6 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => { } } - // If not authenticated with Firebase, try falling back to API key if available return useApiKeyAuthStore().getAuthHeader() } diff --git a/src/stores/workspaceAuthStore.ts b/src/stores/workspaceAuthStore.ts new file mode 100644 index 000000000..33e57fdea --- /dev/null +++ b/src/stores/workspaceAuthStore.ts @@ -0,0 +1,373 @@ +import { defineStore } from 'pinia' +import { computed, ref, shallowRef } from 'vue' +import { z } from 'zod' +import { fromZodError } from 'zod-validation-error' + +import { t } from '@/i18n' +import { + TOKEN_REFRESH_BUFFER_MS, + WORKSPACE_STORAGE_KEYS +} from '@/platform/auth/workspace/workspaceConstants' +import { remoteConfig } from '@/platform/remoteConfig/remoteConfig' +import { api } from '@/scripts/api' +import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore' +import type { AuthHeader } from '@/types/authTypes' +import type { WorkspaceWithRole } from '@/platform/auth/workspace/workspaceTypes' + +const WorkspaceWithRoleSchema = z.object({ + id: z.string(), + name: z.string(), + type: z.enum(['personal', 'team']), + role: z.enum(['owner', 'member']) +}) + +const WorkspaceTokenResponseSchema = z.object({ + token: z.string(), + expires_at: z.string(), + workspace: z.object({ + id: z.string(), + name: z.string(), + type: z.enum(['personal', 'team']) + }), + role: z.enum(['owner', 'member']), + permissions: z.array(z.string()) +}) + +export class WorkspaceAuthError extends Error { + constructor( + message: string, + public readonly code?: string + ) { + super(message) + this.name = 'WorkspaceAuthError' + } +} + +export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => { + // State + const currentWorkspace = shallowRef(null) + const workspaceToken = ref(null) + const isLoading = ref(false) + const error = ref(null) + + // Timer state + let refreshTimerId: ReturnType | null = null + + // Request ID to prevent stale refresh operations from overwriting newer workspace contexts + let refreshRequestId = 0 + + // Getters + const isAuthenticated = computed( + () => currentWorkspace.value !== null && workspaceToken.value !== null + ) + + // Private helpers + function stopRefreshTimer(): void { + if (refreshTimerId !== null) { + clearTimeout(refreshTimerId) + refreshTimerId = null + } + } + + function scheduleTokenRefresh(expiresAt: number): void { + stopRefreshTimer() + const now = Date.now() + const refreshAt = expiresAt - TOKEN_REFRESH_BUFFER_MS + const delay = Math.max(0, refreshAt - now) + + refreshTimerId = setTimeout(() => { + void refreshToken() + }, delay) + } + + function persistToSession( + workspace: WorkspaceWithRole, + token: string, + expiresAt: number + ): void { + try { + sessionStorage.setItem( + WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE, + JSON.stringify(workspace) + ) + sessionStorage.setItem(WORKSPACE_STORAGE_KEYS.TOKEN, token) + sessionStorage.setItem( + WORKSPACE_STORAGE_KEYS.EXPIRES_AT, + expiresAt.toString() + ) + } catch { + console.warn('Failed to persist workspace context to sessionStorage') + } + } + + function clearSessionStorage(): void { + try { + sessionStorage.removeItem(WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE) + sessionStorage.removeItem(WORKSPACE_STORAGE_KEYS.TOKEN) + sessionStorage.removeItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT) + } catch { + console.warn('Failed to clear workspace context from sessionStorage') + } + } + + // Actions + function init(): void { + initializeFromSession() + } + + function destroy(): void { + stopRefreshTimer() + } + + function initializeFromSession(): boolean { + if (!remoteConfig.value.team_workspaces_enabled) { + return false + } + + try { + const workspaceJson = sessionStorage.getItem( + WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE + ) + const token = sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN) + const expiresAtStr = sessionStorage.getItem( + WORKSPACE_STORAGE_KEYS.EXPIRES_AT + ) + + if (!workspaceJson || !token || !expiresAtStr) { + return false + } + + const expiresAt = parseInt(expiresAtStr, 10) + if (isNaN(expiresAt) || expiresAt <= Date.now()) { + clearSessionStorage() + return false + } + + const parsedWorkspace = JSON.parse(workspaceJson) + const parseResult = WorkspaceWithRoleSchema.safeParse(parsedWorkspace) + + if (!parseResult.success) { + clearSessionStorage() + return false + } + + currentWorkspace.value = parseResult.data + workspaceToken.value = token + error.value = null + + scheduleTokenRefresh(expiresAt) + return true + } catch { + clearSessionStorage() + return false + } + } + + async function switchWorkspace(workspaceId: string): Promise { + if (!remoteConfig.value.team_workspaces_enabled) { + return + } + + // Only increment request ID when switching to a different workspace + // This invalidates stale refresh operations for the old workspace + // but allows refresh operations for the same workspace to complete + if (currentWorkspace.value?.id !== workspaceId) { + refreshRequestId++ + } + + isLoading.value = true + error.value = null + + try { + const firebaseAuthStore = useFirebaseAuthStore() + const firebaseToken = await firebaseAuthStore.getIdToken() + if (!firebaseToken) { + throw new WorkspaceAuthError( + t('workspaceAuth.errors.notAuthenticated'), + 'NOT_AUTHENTICATED' + ) + } + + const response = await fetch(api.apiURL('/auth/token'), { + method: 'POST', + headers: { + Authorization: `Bearer ${firebaseToken}`, + 'Content-Type': 'application/json' + }, + body: JSON.stringify({ workspace_id: workspaceId }) + }) + + if (!response.ok) { + const errorData = await response.json().catch(() => ({})) + const message = errorData.message || response.statusText + + if (response.status === 401) { + throw new WorkspaceAuthError( + t('workspaceAuth.errors.invalidFirebaseToken'), + 'INVALID_FIREBASE_TOKEN' + ) + } + if (response.status === 403) { + throw new WorkspaceAuthError( + t('workspaceAuth.errors.accessDenied'), + 'ACCESS_DENIED' + ) + } + if (response.status === 404) { + throw new WorkspaceAuthError( + t('workspaceAuth.errors.workspaceNotFound'), + 'WORKSPACE_NOT_FOUND' + ) + } + + throw new WorkspaceAuthError( + t('workspaceAuth.errors.tokenExchangeFailed', { error: message }), + 'TOKEN_EXCHANGE_FAILED' + ) + } + + const rawData = await response.json() + const parseResult = WorkspaceTokenResponseSchema.safeParse(rawData) + + if (!parseResult.success) { + throw new WorkspaceAuthError( + t('workspaceAuth.errors.tokenExchangeFailed', { + error: fromZodError(parseResult.error).message + }), + 'TOKEN_EXCHANGE_FAILED' + ) + } + + const data = parseResult.data + const expiresAt = new Date(data.expires_at).getTime() + + if (isNaN(expiresAt)) { + throw new WorkspaceAuthError( + t('workspaceAuth.errors.tokenExchangeFailed', { + error: 'Invalid expiry timestamp' + }), + 'TOKEN_EXCHANGE_FAILED' + ) + } + + const workspaceWithRole: WorkspaceWithRole = { + ...data.workspace, + role: data.role + } + + currentWorkspace.value = workspaceWithRole + workspaceToken.value = data.token + + persistToSession(workspaceWithRole, data.token, expiresAt) + scheduleTokenRefresh(expiresAt) + } catch (err) { + error.value = err instanceof Error ? err : new Error(String(err)) + throw error.value + } finally { + isLoading.value = false + } + } + + async function refreshToken(): Promise { + if (!currentWorkspace.value) { + return + } + + const workspaceId = currentWorkspace.value.id + // Capture the current request ID to detect if workspace context changed during refresh + const capturedRequestId = refreshRequestId + const maxRetries = 3 + const baseDelayMs = 1000 + + for (let attempt = 0; attempt <= maxRetries; attempt++) { + // Check if workspace context changed since refresh started (user switched workspaces) + if (capturedRequestId !== refreshRequestId) { + console.warn( + 'Aborting stale token refresh: workspace context changed during refresh' + ) + return + } + + try { + await switchWorkspace(workspaceId) + return + } catch (err) { + const isAuthError = err instanceof WorkspaceAuthError + + const isPermanentError = + isAuthError && + (err.code === 'ACCESS_DENIED' || + err.code === 'WORKSPACE_NOT_FOUND' || + err.code === 'INVALID_FIREBASE_TOKEN' || + err.code === 'NOT_AUTHENTICATED') + + if (isPermanentError) { + // Only clear context if this refresh is still for the current workspace + if (capturedRequestId === refreshRequestId) { + console.error('Workspace access revoked or auth invalid:', err) + clearWorkspaceContext() + } + return + } + + const isTransientError = + isAuthError && err.code === 'TOKEN_EXCHANGE_FAILED' + + if (isTransientError && attempt < maxRetries) { + const delay = baseDelayMs * Math.pow(2, attempt) + console.warn( + `Token refresh failed (attempt ${attempt + 1}/${maxRetries + 1}), retrying in ${delay}ms:`, + err + ) + await new Promise((resolve) => setTimeout(resolve, delay)) + continue + } + + // Only clear context if this refresh is still for the current workspace + if (capturedRequestId === refreshRequestId) { + console.error('Failed to refresh workspace token after retries:', err) + clearWorkspaceContext() + } + } + } + } + + function getWorkspaceAuthHeader(): AuthHeader | null { + if (!workspaceToken.value) { + return null + } + return { + Authorization: `Bearer ${workspaceToken.value}` + } + } + + function clearWorkspaceContext(): void { + // Increment request ID to invalidate any in-flight stale refresh operations + refreshRequestId++ + stopRefreshTimer() + currentWorkspace.value = null + workspaceToken.value = null + error.value = null + clearSessionStorage() + } + + return { + // State + currentWorkspace, + workspaceToken, + isLoading, + error, + + // Getters + isAuthenticated, + + // Actions + init, + destroy, + initializeFromSession, + switchWorkspace, + refreshToken, + getWorkspaceAuthHeader, + clearWorkspaceContext + } +})