From 64319ec2a22c22e34fa2e30d0115ae77dcabcd20 Mon Sep 17 00:00:00 2001
From: Phil Pluckthun <phil@kitten.sh>
Date: Tue, 7 Mar 2023 16:57:29 +0000
Subject: [PATCH] fix(auth): Fix willAuthError not being called on waiting
 operations (#3017)

---
 .changeset/stupid-cobras-clap.md        |  5 ++
 exchanges/auth/src/authExchange.test.ts | 50 +++++++++++++++++
 exchanges/auth/src/authExchange.ts      | 75 +++++++++++--------------
 3 files changed, 87 insertions(+), 43 deletions(-)
 create mode 100644 .changeset/stupid-cobras-clap.md

diff --git a/.changeset/stupid-cobras-clap.md b/.changeset/stupid-cobras-clap.md
new file mode 100644
index 0000000000..b14db43bea
--- /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.test.ts b/exchanges/auth/src/authExchange.test.ts
index 5227688d67..65d7ad0073 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'
+  );
+});
diff --git a/exchanges/auth/src/authExchange.ts b/exchanges/auth/src/authExchange.ts
index a57bdbfae0..ddce6b61c8 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);