Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 39 additions & 8 deletions packages/cli-kit/src/private/node/session.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,6 @@ const validIdentityToken: IdentityToken = {
expiresAt: futureDate,
scopes: ['scope', 'scope2'],
userId,
alias: userId,
}

const validTokens: OAuthSession = {
Expand Down Expand Up @@ -229,7 +228,7 @@ The CLI is currently unable to prompt for reauthentication.`,
expect(fetchSessions).toHaveBeenCalledOnce()
})

test('falls back to userId when email fetch fails', async () => {
test('falls back to userId as alias when email fetch fails', async () => {
// Given
vi.mocked(validateSession).mockResolvedValueOnce('needs_full_auth')
vi.mocked(fetchSessions).mockResolvedValue(undefined)
Expand All @@ -243,7 +242,6 @@ The CLI is currently unable to prompt for reauthentication.`,
expect(businessPlatformRequest).toHaveBeenCalled()
expect(storeSessions).toHaveBeenCalledOnce()

// Verify the session was stored with userId as alias (fallback)
const storedSession = vi.mocked(storeSessions).mock.calls[0]![0]
expect(storedSession[fqdn]![userId]!.identity.alias).toBe(userId)

Expand All @@ -262,7 +260,7 @@ The CLI is currently unable to prompt for reauthentication.`,
vi.mocked(exchangeAccessForApplicationTokens).mockResolvedValueOnce(appTokensWithoutBusinessPlatform)

// When
const got = await ensureAuthenticated(defaultApplications)
await ensureAuthenticated(defaultApplications)

// Then
expect(businessPlatformRequest).not.toHaveBeenCalled()
Expand Down Expand Up @@ -650,16 +648,24 @@ describe('ensureAuthenticated email fetch functionality', () => {

test('preserves existing alias during refresh token flow', async () => {
// Given
const sessionsWithAlias: Sessions = {
[fqdn]: {
[userId]: {
identity: {...validIdentityToken, alias: 'user@example.com'},
applications: appTokens,
},
},
}
vi.mocked(validateSession).mockResolvedValueOnce('needs_refresh')
vi.mocked(fetchSessions).mockResolvedValue(validSessions)
vi.mocked(fetchSessions).mockResolvedValue(sessionsWithAlias)

// When
const got = await ensureAuthenticated(defaultApplications)

// Then
// The email fetch is not called during refresh - the session keeps its existing alias
expect(businessPlatformRequest).not.toHaveBeenCalled()
expect(storeSessions).toBeCalledWith(validSessions)
const storedSession = vi.mocked(storeSessions).mock.calls[0]![0]
expect(storedSession[fqdn]![userId]!.identity.alias).toBe('user@example.com')
expect(got).toEqual(validTokens)
})

Expand All @@ -684,7 +690,32 @@ describe('ensureAuthenticated email fetch functionality', () => {
expect(got).toEqual(validTokens)
})

test('uses userId as alias when email is not available', async () => {
test('preserves existing alias during token refresh error fallback when email fetch fails', async () => {
// Given
const sessionsWithAlias: Sessions = {
[fqdn]: {
[userId]: {
identity: {...validIdentityToken, alias: 'my-custom-alias'},
applications: {},
},
},
}
const tokenResponseError = new InvalidGrantError()
vi.mocked(validateSession).mockResolvedValueOnce('needs_refresh')
vi.mocked(fetchSessions).mockResolvedValue(sessionsWithAlias)
vi.mocked(refreshAccessToken).mockRejectedValueOnce(tokenResponseError)
vi.mocked(businessPlatformRequest).mockRejectedValueOnce(new Error('API Error'))

// When
const got = await ensureAuthenticated(defaultApplications)

// Then
const storedSession = vi.mocked(storeSessions).mock.calls[0]![0]
expect(storedSession[fqdn]![userId]!.identity.alias).toBe('my-custom-alias')
expect(got).toEqual(validTokens)
})

test('falls back to userId as alias when email is not available', async () => {
// Given
vi.mocked(validateSession).mockResolvedValueOnce('needs_full_auth')
vi.mocked(fetchSessions).mockResolvedValue(undefined)
Expand Down
18 changes: 9 additions & 9 deletions packages/cli-kit/src/private/node/session.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ async function fetchEmail(businessPlatformToken: string | undefined): Promise<st

try {
const userEmailResult = await businessPlatformRequest<UserEmailQuery>(UserEmailQueryString, businessPlatformToken)
return userEmailResult.currentUserAccount?.email
return userEmailResult.currentUserAccount?.email ?? undefined
Comment thread
gonzaloriestra marked this conversation as resolved.
// eslint-disable-next-line no-catch-all/no-catch-all
} catch (error) {
outputDebug(outputContent`Failed to fetch user email: ${(error as Error).message ?? String(error)}`)
Expand Down Expand Up @@ -230,15 +230,15 @@ ${outputToken.json(applications)}
if (validationResult === 'needs_full_auth') {
await throwOnNoPrompt(noPrompt)
outputDebug(outputContent`Initiating the full authentication flow...`)
newSession = await executeCompleteFlow(applications)
newSession = await executeCompleteFlow(applications, currentSession?.identity.alias)
Comment thread
gonzaloriestra marked this conversation as resolved.
} else if (validationResult === 'needs_refresh' || forceRefresh) {
outputDebug(outputContent`The current session is valid but needs refresh. Refreshing...`)
try {
newSession = await refreshTokens(currentSession!, applications)
} catch (error) {
if (error instanceof InvalidGrantError) {
await throwOnNoPrompt(noPrompt)
newSession = await executeCompleteFlow(applications)
newSession = await executeCompleteFlow(applications, currentSession?.identity.alias)
} else if (error instanceof InvalidRequestError) {
await sessionStore.remove()
throw new AbortError('\nError validating auth session', "We've cleared the current session, please try again")
Expand Down Expand Up @@ -289,9 +289,9 @@ The CLI is currently unable to prompt for reauthentication.`,
* Execute the full authentication flow.
*
* @param applications - An object containing the applications we need to be authenticated with.
* @param alias - Optional alias to use for the session.
* @param existingAlias - Optional alias from a previous session to preserve if the email fetch fails.
*/
async function executeCompleteFlow(applications: OAuthApplications): Promise<Session> {
async function executeCompleteFlow(applications: OAuthApplications, existingAlias?: string): Promise<Session> {
const scopes = getFlattenScopes(applications)
const exchangeScopes = getExchangeScopes(applications)
const store = applications.adminApi?.storeFqdn
Expand All @@ -318,9 +318,9 @@ async function executeCompleteFlow(applications: OAuthApplications): Promise<Ses
outputDebug(outputContent`CLI token received. Exchanging it for application tokens...`)
const result = await exchangeAccessForApplicationTokens(identityToken, exchangeScopes, store)

// Get the alias for the session (email or userId)
// Preserve existing alias if available, otherwise try fetching email
const businessPlatformToken = result[applicationId('business-platform')]?.accessToken
const alias = (await fetchEmail(businessPlatformToken)) ?? identityToken.userId
const alias = existingAlias ?? (await fetchEmail(businessPlatformToken)) ?? identityToken.userId

const session: Session = {
identity: {
Expand Down Expand Up @@ -352,7 +352,7 @@ async function refreshTokens(session: Session, applications: OAuthApplications):
)

return {
identity: identityToken,
identity: {...identityToken, alias: session.identity.alias},
applications: applicationTokens,
}
}
Expand Down Expand Up @@ -447,6 +447,6 @@ function buildIdentityTokenFromEnv(
...identityTokenInformation,
expiresAt: new Date(Date.now() + 30 * 24 * 60 * 60 * 1000),
scopes,
alias: identityTokenInformation.userId,
alias: undefined,
}
}
51 changes: 49 additions & 2 deletions packages/cli-kit/src/private/node/session/store.test.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import {Sessions} from './schema.js'
import {store, fetch, remove, getSessionAlias, findSessionByAlias} from './store.js'
import {store, fetch, remove, getSessionAlias, setSessionAlias, findSessionByAlias} from './store.js'
import {getSessions, removeSessions, setSessions, removeCurrentSessionId} from '../conf-store.js'
import {identityFqdn} from '../../../public/node/context/fqdn.js'

Expand Down Expand Up @@ -28,7 +28,6 @@ const mockSessions: Sessions = {
expiresAt: new Date(),
scopes: ['scope2'],
userId: 'user2',
alias: 'user2',
},
applications: {},
},
Expand Down Expand Up @@ -132,6 +131,54 @@ describe('session store', () => {
})
})

describe('setSessionAlias', () => {
test('updates alias for an existing user', async () => {
// Given
vi.mocked(getSessions).mockReturnValue(JSON.stringify(mockSessions))

// When
await setSessionAlias('user1', 'New Alias')

// Then
const storedSessions = JSON.parse(vi.mocked(setSessions).mock.calls[0]![0])
expect(storedSessions['identity.fqdn.com'].user1.identity.alias).toBe('New Alias')
})

test('does nothing when no sessions exist', async () => {
// Given
vi.mocked(getSessions).mockReturnValue(undefined)

// When
await setSessionAlias('user1', 'New Alias')

// Then
expect(setSessions).not.toHaveBeenCalled()
})

test('does nothing when user does not exist', async () => {
// Given
vi.mocked(getSessions).mockReturnValue(JSON.stringify(mockSessions))

// When
await setSessionAlias('nonexistent', 'New Alias')

// Then
expect(setSessions).not.toHaveBeenCalled()
})

test('does nothing when fqdn does not match', async () => {
// Given
vi.mocked(getSessions).mockReturnValue(JSON.stringify(mockSessions))
vi.mocked(identityFqdn).mockResolvedValue('different.fqdn.com')

// When
await setSessionAlias('user1', 'New Alias')

// Then
expect(setSessions).not.toHaveBeenCalled()
})
})

describe('findSessionByAlias', () => {
test('returns userId for existing alias', async () => {
// Given
Expand Down
20 changes: 19 additions & 1 deletion packages/cli-kit/src/private/node/session/store.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import {SessionsSchema} from './schema.js'
import {getSessions, removeCurrentSessionId, removeSessions, setSessions} from '../conf-store.js'
import {identityFqdn} from '../../../public/node/context/fqdn.js'
import type {Sessions} from './schema.js'
import type {Session, Sessions} from './schema.js'

/**
* Serializes the session as a JSON and stores it in the system.
Expand Down Expand Up @@ -56,6 +56,24 @@ export async function getSessionAlias(userId: string): Promise<string | undefine
return sessions[fqdn][userId].identity.alias
}

/**
* Sets the alias for a given user's session and persists it.
*
* @param userId - The user ID of the session to update.
* @param alias - The new alias to set.
*/
export async function setSessionAlias(userId: string, alias: string): Promise<void> {
const sessions = await fetch()
if (!sessions) return

const fqdn = await identityFqdn()
const session: Session | undefined = sessions[fqdn]?.[userId]
if (!session) return

session.identity.alias = alias
await store(sessions)
}

/**
* Finds a session by its alias.
*
Expand Down
38 changes: 35 additions & 3 deletions packages/cli-kit/src/public/node/session-prompt.test.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import {promptSessionSelect} from './session-prompt.js'
import {renderSelectPrompt} from './ui.js'
import {renderSelectPrompt, renderTextPrompt} from './ui.js'
import {ensureAuthenticatedUser} from './session.js'
import {identityFqdn} from './context/fqdn.js'
import {setCurrentSessionId} from '../../private/node/conf-store.js'
Expand Down Expand Up @@ -34,8 +34,6 @@ const mockSessions: Sessions = {
expiresAt: new Date(),
scopes: ['scope2'],
userId: 'user2',
// Default alias is same as userId
alias: 'user2',
},
applications: {},
},
Expand Down Expand Up @@ -64,6 +62,21 @@ describe('promptSessionSelect', () => {
expect(result).toEqual('new-alias')
})

test('prompts for alias when no alias is stored', async () => {
// Given
vi.mocked(sessionStore.fetch).mockResolvedValue(undefined)
vi.mocked(sessionStore.getSessionAlias).mockResolvedValue(undefined)
vi.mocked(renderTextPrompt).mockResolvedValue('typed-alias')

// When
const result = await promptSessionSelect()

// Then
expect(renderTextPrompt).toHaveBeenCalled()
expect(sessionStore.setSessionAlias).toHaveBeenCalledWith('new-user-id', 'typed-alias')
expect(result).toEqual('typed-alias')
})

test('prompts user to create new session with alias when no existing sessions', async () => {
// Given
vi.mocked(sessionStore.fetch).mockResolvedValue(undefined)
Expand Down Expand Up @@ -166,6 +179,25 @@ describe('promptSessionSelect', () => {
expect(result).toEqual('custom-alias')
})

test('prompts for alias when selecting "Log in with a different account" and no alias is stored', async () => {
// Given
vi.mocked(sessionStore.fetch).mockResolvedValue(mockSessions)
vi.mocked(renderSelectPrompt).mockResolvedValue('NEW_LOGIN')
vi.mocked(sessionStore.getSessionAlias).mockResolvedValue(undefined)
vi.mocked(renderTextPrompt).mockResolvedValue('my-work-account')

// When
const result = await promptSessionSelect()

// Then
expect(ensureAuthenticatedUser).toHaveBeenCalledWith({}, {forceNewSession: true})
expect(renderTextPrompt).toHaveBeenCalledWith({
message: 'Enter an alias for this account (e.g. your email or a nickname)',
})
expect(sessionStore.setSessionAlias).toHaveBeenCalledWith('new-user-id', 'my-work-account')
expect(result).toEqual('my-work-account')
})

test('does not update alias for existing session when not provided', async () => {
// Given
vi.mocked(sessionStore.fetch).mockResolvedValue(mockSessions)
Expand Down
14 changes: 12 additions & 2 deletions packages/cli-kit/src/public/node/session-prompt.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import {renderSelectPrompt} from './ui.js'
import {renderSelectPrompt, renderTextPrompt} from './ui.js'
import {ensureAuthenticatedUser} from './session.js'
import {identityFqdn} from './context/fqdn.js'
import * as sessionStore from '../../private/node/session/store.js'
Expand Down Expand Up @@ -37,13 +37,23 @@ function buildSessionChoices(sessions: Sessions, fqdn: string): SessionChoice[]

/**
* Handles the new login flow.
* If no alias is stored (email couldn't be fetched), prompts the user for a friendly alias.
*
* @returns The alias of the authenticated user.
*/
async function handleNewLogin(): Promise<string> {
const result = await ensureAuthenticatedUser({}, {forceNewSession: true})
const alias = await sessionStore.getSessionAlias(result.userId)
return alias ?? result.userId

if (!alias) {
const userAlias = await renderTextPrompt({
Comment thread
gonzaloriestra marked this conversation as resolved.
message: 'Enter an alias for this account (e.g. your email or a nickname)',
})
await sessionStore.setSessionAlias(result.userId, userAlias)
return userAlias
}

return alias
}

/**
Expand Down
Loading