Repository: urql-graphql/urql
Pull request #3017 (2023-03-07 16:42:55)

fix(auth): Fix willAuthError not being called on waiting operations
## Summary

When operations are waiting for `authPromise` to resolve, i.e. they're not queued but instead are waiting for the authentication update/refresh to resolve, `willAuthError` will actually _not_ be called, even if it's provided.

This can lead to an easy mistake, where the authentication state may have been rehydrated (on startup), but is never checked to be valid.

We now queue operations that are waiting for the promise like we do for failed operations, which simplifies the logic a little, and make sure to call `willAuthError` on them again.

## Set of changes

- Call `willAuthError` on queued retry operations
- Reuse queue for waiting operations
- Update how `mutate` bypasses the authentication logic



Patch:
diff --git a/.changeset/stupid-cobras-clap.md b/.changeset/stupid-cobras-clap.md
new file mode 100644
index 00000000..b14db43b
--- /dev/null
+++ b/.changeset/stupid-cobras-clap.md
@@ -0,0 +1,5 @@
+---
+'@urql/exchange-auth': patch
+---
+
+Fix `willAuthError` not being called for operations that are waiting on the authentication state to update. This can actually lead to a common issue where operations that came in during the authentication initialization (on startup) will never have `willAuthError` called on them. This can cause an easy mistake where the initial authentication state is never checked to be valid.
diff --git a/exchanges/auth/src/authExchange.ts b/exchanges/auth/src/authExchange.ts
index a57bdbfa..ddce6b61 100644
--- a/exchanges/auth/src/authExchange.ts
+++ b/exchanges/auth/src/authExchange.ts
@@ -1,18 +1,14 @@
 import {
+  Source,
   pipe,
   map,
-  mergeMap,
-  fromPromise,
-  fromValue,
   filter,
   onStart,
-  empty,
   take,
   makeSubject,
   toPromise,
   merge,
   share,
-  takeUntil,
 } from 'wonka';
 
 import {
@@ -72,7 +68,9 @@ export function authExchange<T>({
   willAuthError,
 }: AuthConfig<T>): Exchange {
   return ({ client, forward }) => {
+    const bypassQueue: WeakSet<Operation> = new WeakSet();
     const retryQueue: Map<number, Operation> = new Map();
+
     const {
       source: retrySource$,
       next: retryOperation,
@@ -94,7 +92,10 @@ export function authExchange<T>({
 
         return pipe(
           result$,
-          onStart(() => retryOperation(operation)),
+          onStart(() => {
+            bypassQueue.add(operation);
+            retryOperation(operation);
+          }),
           filter(result => result.operation.key === operation.key),
           take(1),
           toPromise
@@ -142,43 +143,31 @@ export function authExchange<T>({
       );
 
       const opsWithAuth$ = pipe(
-        merge([
-          retrySource$,
-          pipe(
-            pendingOps$,
-            mergeMap(operation => {
-              if (retryQueue.has(operation.key)) {
-                return empty;
-              }
-
-              if (
-                !authPromise &&
-                willAuthError &&
-                willAuthError({ operation, authState })
-              ) {
-                refreshAuth(operation);
-                return empty;
-              } else if (!authPromise) {
-                return fromValue(addAuthAttemptToOperation(operation, false));
-              }
-
-              const teardown$ = pipe(
-                sharedOps$,
-                filter(op => {
-                  return op.kind === 'teardown' && op.key === operation.key;
-                })
-              );
-
-              return pipe(
-                fromPromise(authPromise),
-                map(() => addAuthAttemptToOperation(operation, false)),
-                takeUntil(teardown$)
-              );
-            })
-          ),
-        ]),
-        map(operation => addAuthToOperation({ operation, authState }))
-      );
+        merge([retrySource$, pendingOps$]),
+        map(operation => {
+          if (bypassQueue.has(operation)) {
+            return operation;
+          } else if (authPromise) {
+            operation = addAuthAttemptToOperation(operation, false);
+            retryQueue.set(
+              operation.key,
+              addAuthAttemptToOperation(operation, false)
+            );
+            return null;
+          } else if (
+            !operation.context.authAttempt &&
+            willAuthError &&
+            willAuthError({ operation, authState })
+          ) {
+            refreshAuth(operation);
+            return null;
+          }
+
+          operation = addAuthAttemptToOperation(operation, false);
+          return addAuthToOperation({ operation, authState });
+        }),
+        filter(Boolean)
+      ) as Source<Operation>;
 
       const result$ = pipe(merge([opsWithAuth$, teardownOps$]), forward, share);
 


Test patch:
diff --git a/exchanges/auth/src/authExchange.test.ts b/exchanges/auth/src/authExchange.test.ts
index 5227688d..65d7ad00 100644
--- a/exchanges/auth/src/authExchange.test.ts
+++ b/exchanges/auth/src/authExchange.test.ts
@@ -165,6 +165,8 @@ it('adds the same token to subsequent operations', async () => {
     toPromise
   );
 
+  await new Promise(resolve => setTimeout(resolve));
+
   next(queryOperation);
 
   next(
@@ -333,3 +335,51 @@ it('triggers authentication when an operation will error', async () => {
     'final-token'
   );
 });
+
+it('calls willAuthError on queued operations', async () => {
+  const { exchangeArgs, result, operations } = makeExchangeArgs();
+  const { source, next } = makeSubject<any>();
+
+  let initialAuthResolve: ((_?: any) => void) | undefined;
+
+  const willAuthError = vi.fn().mockReturnValue(false);
+
+  pipe(
+    source,
+    authExchange<{ token: string }>({
+      async getAuth() {
+        await new Promise(resolve => {
+          initialAuthResolve = resolve;
+        });
+
+        return { token: 'token' };
+      },
+      willAuthError,
+      didAuthError: () => false,
+      addAuthToOperation: ({ authState, operation }) => {
+        return withAuthHeader(operation, authState?.token);
+      },
+    })(exchangeArgs),
+    publish
+  );
+
+  await Promise.resolve();
+
+  next(queryOperation);
+  expect(result).toHaveBeenCalledTimes(0);
+  expect(willAuthError).toHaveBeenCalledTimes(0);
+
+  expect(initialAuthResolve).toBeDefined();
+  initialAuthResolve!();
+
+  await new Promise(resolve => setTimeout(resolve));
+
+  expect(willAuthError).toHaveBeenCalledTimes(1);
+  expect(result).toHaveBeenCalledTimes(1);
+
+  expect(operations.length).toBe(1);
+  expect(operations[0]).toHaveProperty(
+    'context.fetchOptions.headers.Authorization',
+    'token'
+  );
+});