Revert "Simplify authentication flow"

Author: isaacroldan

packages/cli-kit/src/private/node/session.test.ts

@@ -7,11 +7,17 @@ import {
   setLastSeenAuthMethod,
   setLastSeenUserIdAfterAuth,
 } from './session.js'
-import {exchangeCustomPartnerToken, refreshAccessToken, InvalidGrantError} from './session/exchange.js'
+import {
+  exchangeAccessForApplicationTokens,
+  exchangeCustomPartnerToken,
+  refreshAccessToken,
+  InvalidGrantError,
+} from './session/exchange.js'
 import {allDefaultScopes} from './session/scopes.js'
 import {store as storeSessions, fetch as fetchSessions, remove as secureRemove} from './session/store.js'
-import {IdentityToken, Sessions} from './session/schema.js'
+import {ApplicationToken, IdentityToken, Sessions} from './session/schema.js'
 import {validateSession} from './session/validate.js'
+import {applicationId} from './session/identity.js'
 import {pollForDeviceAuthorization, requestDeviceAuthorization} from './session/device-authorization.js'
 import {getCurrentSessionId} from './conf-store.js'
 import * as fqdnModule from '../../public/node/context/fqdn.js'
@@ -44,15 +50,54 @@ const validIdentityToken: IdentityToken = {
 }
 
 const validTokens: OAuthSession = {
-  admin: {token: 'access_token', storeFqdn: 'mystore.myshopify.com'},
-  storefront: 'access_token',
-  partners: 'access_token',
+  admin: {token: 'admin_token', storeFqdn: 'mystore.myshopify.com'},
+  storefront: 'storefront_token',
+  partners: 'partners_token',
   userId,
 }
 
+const appTokens: Record<string, ApplicationToken> = {
+  // Admin APIs includes domain in the key
+  'mystore.myshopify.com-admin': {
+    accessToken: 'admin_token',
+    expiresAt: futureDate,
+    scopes: ['scope', 'scope2'],
+  },
+  'storefront-renderer': {
+    accessToken: 'storefront_token',
+    expiresAt: futureDate,
+    scopes: ['scope1'],
+  },
+  partners: {
+    accessToken: 'partners_token',
+    expiresAt: futureDate,
+    scopes: ['scope2'],
+  },
+  'business-platform': {
+    accessToken: 'business_platform_token',
+    expiresAt: futureDate,
+    scopes: ['scope3'],
+  },
+}
+
+const partnersToken: ApplicationToken = {
+  accessToken: 'custom_partners_token',
+  expiresAt: futureDate,
+  scopes: ['scope2'],
+}
+
 const fqdn = 'fqdn.com'
 
 const validSessions: Sessions = {
+  [fqdn]: {
+    [userId]: {
+      identity: validIdentityToken,
+      applications: appTokens,
+    },
+  },
+}
+
+const invalidSessions: Sessions = {
   [fqdn]: {
     [userId]: {
       identity: validIdentityToken,
@@ -62,6 +107,7 @@ const validSessions: Sessions = {
 }
 
 vi.mock('../../public/node/context/local.js')
+vi.mock('./session/identity')
 vi.mock('./session/authorize')
 vi.mock('./session/exchange')
 vi.mock('./session/scopes')
@@ -77,9 +123,11 @@ vi.mock('../../public/node/system.js')
 
 beforeEach(() => {
   vi.spyOn(fqdnModule, 'identityFqdn').mockResolvedValue(fqdn)
+  vi.mocked(exchangeAccessForApplicationTokens).mockResolvedValue(appTokens)
   vi.mocked(refreshAccessToken).mockResolvedValue(validIdentityToken)
+  vi.mocked(applicationId).mockImplementation((app) => app)
   vi.mocked(exchangeCustomPartnerToken).mockResolvedValue({
-    accessToken: 'custom_partners_token',
+    accessToken: partnersToken.accessToken,
     userId: validIdentityToken.userId,
   })
   vi.mocked(partnersRequest).mockResolvedValue(undefined)
@@ -114,6 +162,7 @@ describe('ensureAuthenticated when previous session is invalid', () => {
     const got = await ensureAuthenticated(defaultApplications)
 
     // Then
+    expect(exchangeAccessForApplicationTokens).toBeCalled()
     expect(refreshAccessToken).not.toBeCalled()
     expect(businessPlatformRequest).toHaveBeenCalled()
     expect(storeSessions).toHaveBeenCalledOnce()
@@ -153,16 +202,16 @@ The CLI is currently unable to prompt for reauthentication.`,
   test('executes complete auth flow if session is for a different fqdn', async () => {
     // Given
     vi.mocked(validateSession).mockResolvedValueOnce('needs_full_auth')
-    vi.mocked(fetchSessions).mockResolvedValue(validSessions)
+    vi.mocked(fetchSessions).mockResolvedValue(invalidSessions)
     const expectedSessions = {
-      ...validSessions,
+      ...invalidSessions,
       [fqdn]: {
         [userId]: {
           identity: {
             ...validIdentityToken,
             alias: 'user@example.com',
           },
-          applications: {},
+          applications: appTokens,
         },
       },
     }
@@ -171,7 +220,7 @@ The CLI is currently unable to prompt for reauthentication.`,
     const got = await ensureAuthenticated(defaultApplications)
 
     // Then
-
+    expect(exchangeAccessForApplicationTokens).toBeCalled()
     expect(refreshAccessToken).not.toBeCalled()
     expect(storeSessions).toBeCalledWith(expectedSessions)
     expect(got).toEqual(validTokens)
@@ -190,7 +239,7 @@ The CLI is currently unable to prompt for reauthentication.`,
     const got = await ensureAuthenticated(defaultApplications)
 
     // Then
-
+    expect(exchangeAccessForApplicationTokens).toBeCalled()
     expect(businessPlatformRequest).toHaveBeenCalled()
     expect(storeSessions).toHaveBeenCalledOnce()
 
@@ -201,20 +250,26 @@ The CLI is currently unable to prompt for reauthentication.`,
     expect(got).toEqual(validTokens)
   })
 
-  test('uses identity token to fetch email during full auth flow', async () => {
+  test('falls back to userId when no business platform token available', async () => {
     // Given
     vi.mocked(validateSession).mockResolvedValueOnce('needs_full_auth')
     vi.mocked(fetchSessions).mockResolvedValue(undefined)
+    const appTokensWithoutBusinessPlatform = {
+      'mystore.myshopify.com-admin': appTokens['mystore.myshopify.com-admin']!,
+      'storefront-renderer': appTokens['storefront-renderer']!,
+      partners: appTokens.partners!,
+    }
+    vi.mocked(exchangeAccessForApplicationTokens).mockResolvedValueOnce(appTokensWithoutBusinessPlatform)
 
     // When
     const got = await ensureAuthenticated(defaultApplications)
 
     // Then
-    expect(businessPlatformRequest).toHaveBeenCalledWith(expect.any(String), 'access_token')
+    expect(businessPlatformRequest).not.toHaveBeenCalled()
 
-    // Verify the session was stored with email as alias
+    // 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('user@example.com')
+    expect(storedSession[fqdn]![userId]!.identity.alias).toBe(userId)
   })
 
   test('executes complete auth flow if requesting additional scopes', async () => {
@@ -226,7 +281,7 @@ The CLI is currently unable to prompt for reauthentication.`,
     const got = await ensureAuthenticated(defaultApplications)
 
     // Then
-
+    expect(exchangeAccessForApplicationTokens).toBeCalled()
     expect(refreshAccessToken).not.toBeCalled()
     expect(businessPlatformRequest).toHaveBeenCalled()
     expect(storeSessions).toHaveBeenCalledOnce()
@@ -252,7 +307,7 @@ describe('when existing session is valid', () => {
     const got = await ensureAuthenticated(defaultApplications)
 
     // Then
-
+    expect(exchangeAccessForApplicationTokens).not.toBeCalled()
     expect(refreshAccessToken).not.toBeCalled()
     expect(got).toEqual(validTokens)
     await expect(getLastSeenUserIdAfterAuth()).resolves.toBe('1234-5678')
@@ -265,18 +320,13 @@ describe('when existing session is valid', () => {
     vi.mocked(validateSession).mockResolvedValueOnce('ok')
     vi.mocked(fetchSessions).mockResolvedValue(validSessions)
     vi.mocked(getPartnersToken).mockReturnValue('custom_cli_token')
-    const expected = {
-      admin: {token: 'access_token', storeFqdn: 'mystore.myshopify.com'},
-      storefront: 'access_token',
-      partners: 'custom_partners_token',
-      userId,
-    }
+    const expected = {...validTokens, partners: 'custom_partners_token'}
 
     // When
     const got = await ensureAuthenticated(defaultApplications)
 
     // Then
-
+    expect(exchangeAccessForApplicationTokens).not.toBeCalled()
     expect(refreshAccessToken).not.toBeCalled()
     expect(got).toEqual(expected)
     await expect(getLastSeenUserIdAfterAuth()).resolves.toBe('1234-5678')
@@ -294,16 +344,8 @@ describe('when existing session is valid', () => {
 
     // Then
     expect(refreshAccessToken).toBeCalled()
-
-    const expectedSessions = {
-      [fqdn]: {
-        [userId]: {
-          identity: validIdentityToken,
-          applications: {},
-        },
-      },
-    }
-    expect(storeSessions).toBeCalledWith(expectedSessions)
+    expect(exchangeAccessForApplicationTokens).toBeCalled()
+    expect(storeSessions).toBeCalledWith(validSessions)
     expect(got).toEqual(validTokens)
     await expect(getLastSeenUserIdAfterAuth()).resolves.toBe('1234-5678')
     await expect(getLastSeenAuthMethod()).resolves.toEqual('device_auth')
@@ -322,16 +364,8 @@ describe('when existing session is expired', () => {
 
     // Then
     expect(refreshAccessToken).toBeCalled()
-
-    const expectedSessions = {
-      [fqdn]: {
-        [userId]: {
-          identity: validIdentityToken,
-          applications: {},
-        },
-      },
-    }
-    expect(storeSessions).toBeCalledWith(expectedSessions)
+    expect(exchangeAccessForApplicationTokens).toBeCalled()
+    expect(storeSessions).toBeCalledWith(validSessions)
     expect(got).toEqual(validTokens)
     await expect(getLastSeenUserIdAfterAuth()).resolves.toBe('1234-5678')
     await expect(getLastSeenAuthMethod()).resolves.toEqual('device_auth')
@@ -351,7 +385,7 @@ describe('when existing session is expired', () => {
 
     // Then
     expect(refreshAccessToken).toBeCalled()
-
+    expect(exchangeAccessForApplicationTokens).toBeCalled()
     expect(businessPlatformRequest).toHaveBeenCalled()
     expect(storeSessions).toHaveBeenCalledOnce()
 
@@ -610,15 +644,7 @@ describe('ensureAuthenticated email fetch functionality', () => {
     const got = await ensureAuthenticated(defaultApplications)
 
     // Then
-    const expectedSessions = {
-      [fqdn]: {
-        [userId]: {
-          identity: validIdentityToken,
-          applications: {},
-        },
-      },
-    }
-    expect(storeSessions).toBeCalledWith(expectedSessions)
+    expect(storeSessions).toBeCalledWith(validSessions)
     expect(got).toEqual(validTokens)
   })
 
@@ -633,15 +659,7 @@ describe('ensureAuthenticated email fetch functionality', () => {
     // Then
     // The email fetch is not called during refresh - the session keeps its existing alias
     expect(businessPlatformRequest).not.toHaveBeenCalled()
-    const expectedSessions = {
-      [fqdn]: {
-        [userId]: {
-          identity: validIdentityToken,
-          applications: {},
-        },
-      },
-    }
-    expect(storeSessions).toBeCalledWith(expectedSessions)
+    expect(storeSessions).toBeCalledWith(validSessions)
     expect(got).toEqual(validTokens)
   })