Refresh session on 401 for graphQL requests

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

@@ -55,6 +55,7 @@ describe('retryAwareRequest', () => {
         url: 'https://example.com',
       },
       undefined,
+      undefined,
       {
         defaultDelayMs: 500,
         scheduleDelay: mockScheduleDelayFn,
@@ -100,6 +101,7 @@ describe('retryAwareRequest', () => {
           url: 'https://example.com',
         },
         undefined,
+        undefined,
         {
           limitRetriesTo: 7,
           scheduleDelay: mockScheduleDelayFn,
@@ -110,4 +112,54 @@ describe('retryAwareRequest', () => {
     expect(mockRequestFn).toHaveBeenCalledTimes(8)
     expect(mockScheduleDelayFn).toHaveBeenCalledTimes(7)
   })
+
+  test('calls unauthorizedHandler when receiving 401', async () => {
+    const unauthorizedResponse = {
+      status: 401,
+      errors: [
+        {
+          extensions: {
+            code: '401',
+          },
+        } as any,
+      ],
+      headers: new Headers(),
+    }
+
+    const mockRequestFn = vi
+      .fn()
+      .mockImplementationOnce(() => {
+        throw new ClientError(unauthorizedResponse, {query: ''})
+      })
+      .mockImplementationOnce(() => {
+        return Promise.resolve({
+          status: 200,
+          data: {hello: 'world!'},
+          headers: new Headers(),
+        })
+      })
+
+    const mockUnauthorizedHandler = vi.fn()
+
+    await expect(
+      retryAwareRequest(
+        {
+          request: mockRequestFn,
+          url: 'https://example.com',
+        },
+        undefined,
+        mockUnauthorizedHandler,
+        {
+          scheduleDelay: vi.fn((fn) => fn()),
+        },
+      ),
+    ).resolves.toEqual({
+      headers: expect.anything(),
+      status: 200,
+      data: {hello: 'world!'},
+    })
+
+    expect(mockRequestFn).toHaveBeenCalledTimes(2)
+    expect(mockUnauthorizedHandler).toHaveBeenCalledTimes(1)
+  })
 })

packages/cli-kit/src/private/node/api.ts

@@ -40,6 +40,7 @@ type VerboseResponse<T> = {
   | {status: 'client-error'; clientError: ClientError}
   | {status: 'unknown-error'; error: unknown}
   | {status: 'can-retry'; clientError: ClientError; delayMs: number | undefined}
+  | {status: 'unauthorized'; clientError: ClientError; delayMs: number | undefined}
 )
 
 async function makeVerboseRequest<T extends {headers: Headers; status: number}>({
@@ -88,6 +89,16 @@ async function makeVerboseRequest<T extends {headers: Headers; status: number}>(
           requestId: responseHeaders['x-request-id'],
           delayMs,
         }
+      } else if (err.response.status === 401) {
+        return {
+          status: 'unauthorized',
+          clientError: err,
+          duration,
+          sanitizedHeaders,
+          sanitizedUrl,
+          requestId: responseHeaders['x-request-id'],
+          delayMs: 500,
+        }
       }
 
       return {
@@ -169,12 +180,20 @@ ${result.sanitizedHeaders}
         throw result.clientError
       }
     }
+    case 'unauthorized': {
+      if (errorHandler) {
+        throw errorHandler(result.clientError, result.requestId)
+      } else {
+        throw result.clientError
+      }
+    }
   }
 }
 
 export async function retryAwareRequest<T extends {headers: Headers; status: number}>(
   {request, url}: RequestOptions<T>,
   errorHandler?: (error: unknown, requestId: string | undefined) => unknown,
+  unauthorizedHandler?: () => Promise<void>,
   retryOptions: {
     limitRetriesTo?: number
     defaultDelayMs?: number
@@ -211,6 +230,13 @@ ${result.sanitizedHeaders}
       } else {
         throw result.error
       }
+    } else if (result.status === 'unauthorized') {
+      if (unauthorizedHandler) {
+        // eslint-disable-next-line no-await-in-loop
+        await unauthorizedHandler()
+      } else {
+        throw result.clientError
+      }
     }
 
     if (limitRetriesTo <= retriesUsed) {

packages/cli-kit/src/public/node/api/admin.ts

@@ -63,7 +63,17 @@ export async function adminRequestDoc<TResult, TVariables extends Variables>(
     token: session.token,
     addedHeaders,
   }
-  const result = graphqlRequestDoc<TResult, TVariables>({...opts, query, variables, responseOptions})
+  let unauthorizedHandler
+  if ('refresh' in session) {
+    unauthorizedHandler = session.refresh as () => Promise<void>
+  }
+  const result = graphqlRequestDoc<TResult, TVariables>({
+    ...opts,
+    query,
+    variables,
+    responseOptions,
+    unauthorizedHandler,
+  })
   return result
 }
 

packages/cli-kit/src/public/node/api/graphql.test.ts

@@ -54,7 +54,7 @@ describe('graphqlRequest', () => {
       request: expect.any(Function),
       url: mockedAddress,
     }
-    expect(retryAwareRequest).toHaveBeenCalledWith(receivedObject, expect.any(Function))
+    expect(retryAwareRequest).toHaveBeenCalledWith(receivedObject, expect.any(Function), undefined)
   })
 })
 
@@ -95,6 +95,7 @@ describe('graphqlRequestDoc', () => {
         url: mockedAddress,
       },
       expect.any(Function),
+      undefined,
     )
     expect(debugRequest.debugLogRequestInfo).toHaveBeenCalledWith(
       'mockApi',

packages/cli-kit/src/public/node/api/graphql.ts

@@ -26,16 +26,19 @@ interface GraphQLRequestBaseOptions<TResult> {
 type PerformGraphQLRequestOptions<TResult> = GraphQLRequestBaseOptions<TResult> & {
   queryAsString: string
   variables?: Variables
+  unauthorizedHandler?: () => Promise<void>
 }
 
 export type GraphQLRequestOptions<T> = GraphQLRequestBaseOptions<T> & {
   query: RequestDocument
   variables?: Variables
+  unauthorizedHandler?: () => Promise<void>
 }
 
 export type GraphQLRequestDocOptions<TResult, TVariables> = GraphQLRequestBaseOptions<TResult> & {
   query: TypedDocumentNode<TResult, TVariables> | TypedDocumentNode<TResult, Exact<{[key: string]: never}>>
   variables?: TVariables
+  unauthorizedHandler?: () => Promise<void>
 }
 
 export interface GraphQLResponseOptions<T> {
@@ -49,7 +52,7 @@ export interface GraphQLResponseOptions<T> {
  * @param options - GraphQL request options.
  */
 async function performGraphQLRequest<TResult>(options: PerformGraphQLRequestOptions<TResult>) {
-  const {token, addedHeaders, queryAsString, variables, api, url, responseOptions} = options
+  const {token, addedHeaders, queryAsString, variables, api, url, responseOptions, unauthorizedHandler} = options
   const headers = {
     ...addedHeaders,
     ...buildHeaders(token),
@@ -63,6 +66,7 @@ async function performGraphQLRequest<TResult>(options: PerformGraphQLRequestOpti
     const response = await retryAwareRequest(
       {request: () => client.rawRequest<TResult>(queryAsString, variables), url},
       responseOptions?.handleErrors === false ? undefined : errorHandler(api),
+      unauthorizedHandler,
     )
 
     if (responseOptions?.onResponse) {

packages/cli-kit/src/public/node/themes/api.test.ts

@@ -370,31 +370,6 @@ describe('request errors', () => {
       // Then
     }).rejects.toThrowError(AbortError)
   })
-
-  test(`refresh the session when 401 errors happen`, async () => {
-    // Given
-    const id = 123
-    const assets: AssetParams[] = []
-
-    vi.spyOn(session, 'refresh').mockImplementation(vi.fn())
-    vi.mocked(restRequest)
-      .mockResolvedValueOnce({
-        json: {},
-        status: 401,
-        headers: {},
-      })
-      .mockResolvedValueOnce({
-        json: {},
-        status: 207,
-        headers: {},
-      })
-
-    // When
-    await bulkUploadThemeAssets(id, assets, session)
-
-    // Then
-    expect(session.refresh).toHaveBeenCalledOnce()
-  })
 })
 
 describe('bulkUploadThemeAssets', async () => {