Can't get wsLink's retryDelayMs to work

I'm trying to add some backoff for when connecting to my websocket server fails. I've got the following code, but the retryDelayMs function never seems to be called (I don't see any console messages), and the wsLink keeps trying to reconnect immediately:

const TrpcProviderInner: React.FC<{ children: React.ReactNode }> = p => {
  const [queryClient] = useState(() => new QueryClient())
  const [trpcClient] = useState(() =>
    api.createClient({
      transformer: superjson,
      links: [
        loggerLink({
          enabled: opts =>
            process.env.NODE_ENV === 'development' ||
            (opts.direction === 'down' && opts.result instanceof Error),
        }),
        wsLink<Router>({
          client: createWSClient({
            url: getWsUrl(),
            retryDelayMs: (attemptIndex: number) => {
              console.log('retrying', attemptIndex)
              switch (attemptIndex) {
                case 0:
                  return 50
                case 1:
                  return 100
                case 2:
                  return 200
                case 3:
                  return 500
                case 4:
                  return 1000
                default:
                  return 5000
              }
            },
          }),
        }),
      ],
    }),
  )
  return (
    <api.Provider client={trpcClient} queryClient={queryClient}>
      <QueryClientProvider client={queryClient}>{p.children}</QueryClientProvider>
    </api.Provider>
  )
}

The issue arises from the way wsLink resets the connectAttempt counter upon a successful connection. This make total sense. However, in my case, my WebSocket server validates a session cookie in the connect event handler, and closes the connection on an invalid session. The connectAttempt counter is still reset, because this is a successful connection followed immediately by a close().

Maybe it would make sense to only reset the connectAttempt counter after a successful request/response from the WebSocket server.