Use the right token after refresh for Business Platfrom

Author: gonzaloriestra

packages/app/src/cli/utilities/developer-platform-client.test.ts

@@ -2,7 +2,8 @@ import {createUnauthorizedHandler, DeveloperPlatformClient} from './developer-pl
 import {describe, expect, test, vi} from 'vitest'
 
 describe('createUnauthorizedHandler', () => {
-  const mockToken = 'mock-refreshed-token'
+  const mockAppManagementToken = 'mock-app-management-token'
+  const mockBusinessPlatformToken = 'mock-business-platform-token'
   const createMockClient = () => {
     let tokenRefreshPromise: Promise<string> | undefined
     return {
@@ -13,17 +14,21 @@ describe('createUnauthorizedHandler', () => {
       clearCurrentlyRefreshingToken: () => {
         tokenRefreshPromise = undefined
       },
-      unsafeRefreshToken: vi.fn().mockResolvedValue(mockToken),
+      unsafeRefreshToken: vi.fn().mockResolvedValue(mockAppManagementToken),
+      session: vi.fn().mockResolvedValue({
+        token: mockAppManagementToken,
+        businessPlatformToken: mockBusinessPlatformToken,
+      }),
     } as unknown as DeveloperPlatformClient
   }
 
-  test('refreshes token on first attempt', async () => {
+  test('refreshes token on first attempt and returns appManagement token by default', async () => {
     const mockClient = createMockClient()
     const handler = createUnauthorizedHandler(mockClient)
 
     const result = await handler.handler()
 
-    expect(result).toEqual({token: mockToken})
+    expect(result).toEqual({token: mockAppManagementToken})
     expect(mockClient.unsafeRefreshToken).toHaveBeenCalledTimes(1)
   })
 
@@ -35,7 +40,7 @@ describe('createUnauthorizedHandler', () => {
     handler.handler()
     const result = await handler.handler()
 
-    expect(result).toEqual({token: mockToken})
+    expect(result).toEqual({token: mockAppManagementToken})
     expect(mockClient.unsafeRefreshToken).toHaveBeenCalledTimes(1)
   })
 
@@ -51,4 +56,24 @@ describe('createUnauthorizedHandler', () => {
     // The following is evidence that the finally block is called correctly
     expect(mockClient.unsafeRefreshToken).toHaveBeenCalledTimes(2)
   })
+
+  test('returns businessPlatform token when tokenType is businessPlatform', async () => {
+    const mockClient = createMockClient()
+    const handler = createUnauthorizedHandler(mockClient, 'businessPlatform')
+
+    const result = await handler.handler()
+
+    expect(result).toEqual({token: mockBusinessPlatformToken})
+    expect(mockClient.unsafeRefreshToken).toHaveBeenCalledTimes(1)
+  })
+
+  test('returns appManagement token when tokenType is appManagement', async () => {
+    const mockClient = createMockClient()
+    const handler = createUnauthorizedHandler(mockClient, 'appManagement')
+
+    const result = await handler.handler()
+
+    expect(result).toEqual({token: mockAppManagementToken})
+    expect(mockClient.unsafeRefreshToken).toHaveBeenCalledTimes(1)
+  })
 })

packages/app/src/cli/utilities/developer-platform-client.ts

@@ -296,24 +296,30 @@ export interface DeveloperPlatformClient {
 
 const inProgressRefreshes = new WeakMap<DeveloperPlatformClient, Promise<string>>()
 
-export function createUnauthorizedHandler(client: DeveloperPlatformClient): UnauthorizedHandler {
+export function createUnauthorizedHandler(
+  client: DeveloperPlatformClient,
+  tokenType: 'appManagement' | 'businessPlatform' = 'appManagement',
+): UnauthorizedHandler {
   return {
     type: 'token_refresh',
     handler: async () => {
       let tokenRefresher = inProgressRefreshes.get(client)
       if (tokenRefresher) {
-        const token = await tokenRefresher
-        return {token}
+        await tokenRefresher
       } else {
         try {
           tokenRefresher = client.unsafeRefreshToken()
           inProgressRefreshes.set(client, tokenRefresher)
-          const token = await tokenRefresher
-          return {token}
+          await tokenRefresher
         } finally {
           inProgressRefreshes.delete(client)
         }
       }
+
+      // After refresh, get the appropriate token based on the request type
+      const session = await client.session()
+      const token = tokenType === 'businessPlatform' ? session.businessPlatformToken : session.token
+      return {token}
     },
   }
 }

packages/app/src/cli/utilities/developer-platform-client/app-management-client.ts

@@ -273,7 +273,7 @@ export class AppManagementClient implements DeveloperPlatformClient {
           cacheExtraKey: userId,
         },
         token: businessPlatformToken,
-        unauthorizedHandler: this.createUnauthorizedHandler(),
+        unauthorizedHandler: this.createUnauthorizedHandler('businessPlatform'),
       })
 
       if (getPartnersToken() && userInfoResult.currentUserAccount) {
@@ -856,7 +856,7 @@ export class AppManagementClient implements DeveloperPlatformClient {
       token: await this.businessPlatformToken(),
       organizationId: String(numberFromGid(orgId)),
       variables,
-      unauthorizedHandler: this.createUnauthorizedHandler(),
+      unauthorizedHandler: this.createUnauthorizedHandler('appManagement'),
     })
     const provisionResult = fullResult.organizationUserProvisionShopAccess
     if (!provisionResult.success) {
@@ -1057,7 +1057,7 @@ export class AppManagementClient implements DeveloperPlatformClient {
       token: await this.businessPlatformToken(),
       organizationId,
       variables,
-      unauthorizedHandler: this.createUnauthorizedHandler(),
+      unauthorizedHandler: this.createUnauthorizedHandler('appManagement'),
     })
     const result: {[flag: (typeof allBetaFlags)[number]]: boolean} = {}
     allBetaFlags.forEach((flag) => {
@@ -1072,7 +1072,7 @@ export class AppManagementClient implements DeveloperPlatformClient {
     return appManagementRequestDoc({
       ...options,
       token: await this.token(),
-      unauthorizedHandler: this.createUnauthorizedHandler(),
+      unauthorizedHandler: this.createUnauthorizedHandler('appManagement'),
     })
   }
 
@@ -1082,7 +1082,7 @@ export class AppManagementClient implements DeveloperPlatformClient {
     return appDevRequestDoc({
       ...options,
       token: await this.token(),
-      unauthorizedHandler: this.createUnauthorizedHandler(),
+      unauthorizedHandler: this.createUnauthorizedHandler('appManagement'),
     })
   }
 
@@ -1092,7 +1092,7 @@ export class AppManagementClient implements DeveloperPlatformClient {
     return businessPlatformRequestDoc({
       ...options,
       token: await this.businessPlatformToken(),
-      unauthorizedHandler: this.createUnauthorizedHandler(),
+      unauthorizedHandler: this.createUnauthorizedHandler('businessPlatform'),
     })
   }
 
@@ -1102,7 +1102,7 @@ export class AppManagementClient implements DeveloperPlatformClient {
     return businessPlatformOrganizationsRequestDoc({
       ...options,
       token: await this.businessPlatformToken(),
-      unauthorizedHandler: this.createUnauthorizedHandler(),
+      unauthorizedHandler: this.createUnauthorizedHandler('businessPlatform'),
     })
   }
 
@@ -1112,7 +1112,7 @@ export class AppManagementClient implements DeveloperPlatformClient {
     return functionsRequestDoc<TResult, TVariables>({
       ...options,
       token: await this.token(),
-      unauthorizedHandler: this.createUnauthorizedHandler(),
+      unauthorizedHandler: this.createUnauthorizedHandler('appManagement'),
     })
   }
 
@@ -1122,12 +1122,14 @@ export class AppManagementClient implements DeveloperPlatformClient {
     return webhooksRequestDoc<TResult, TVariables>({
       ...options,
       token: await this.token(),
-      unauthorizedHandler: this.createUnauthorizedHandler(),
+      unauthorizedHandler: this.createUnauthorizedHandler('appManagement'),
     })
   }
 
-  private createUnauthorizedHandler(): UnauthorizedHandler {
-    return createUnauthorizedHandler(this)
+  private createUnauthorizedHandler(
+    tokenType: 'appManagement' | 'businessPlatform' = 'appManagement',
+  ): UnauthorizedHandler {
+    return createUnauthorizedHandler(this, tokenType)
   }
 }