fix(entrykit): sign session wallet after top-up (#3697)

Co-authored-by: alvarius <alvarius@lattice.xyz>

Author: karooolis

.changeset/tidy-pillows-learn.md

@@ -0,0 +1,5 @@
+---
+"@latticexyz/entrykit": patch
+---
+
+The login flow now only attempts to register the session account after it has been successfully funded.

packages/entrykit/src/onboarding/ConnectedSteps.tsx

@@ -49,7 +49,7 @@ export function ConnectedSteps({ userClient, initialUserAddress }: Props) {
     }
   }, [closeAccountModal, isNewConnection, prerequisites]);
 
-  const { sessionAddress, hasAllowance, isSpender, hasDelegation, hasGasBalance, hasQuarryBalance } =
+  const { sessionAddress, hasAllowance, isSpender, hasDelegation, hasGasBalance, hasQuarryGasBalance } =
     prerequisites ?? {};
 
   const steps = useMemo((): readonly Step[] => {
@@ -89,7 +89,7 @@ export function ConnectedSteps({ userClient, initialUserAddress }: Props) {
       } else {
         steps.push({
           id: "gasBalanceQuarry",
-          isComplete: !!hasQuarryBalance,
+          isComplete: !!hasQuarryGasBalance,
           content: (props) => <GasBalanceQuarry {...props} userAddress={userAddress} />,
         });
       }
@@ -108,7 +108,7 @@ export function ConnectedSteps({ userClient, initialUserAddress }: Props) {
     hasAllowance,
     hasDelegation,
     hasGasBalance,
-    hasQuarryBalance,
+    hasQuarryGasBalance,
     isSpender,
     paymaster,
     sessionAddress,

packages/entrykit/src/onboarding/Session.tsx

@@ -1,30 +1,41 @@
+import { useEffect } from "react";
+import { Hex } from "viem";
 import { Button } from "../ui/Button";
 import { useSetupSession } from "./useSetupSession";
 import { ConnectedClient } from "../common";
-import { useEffect } from "react";
 import { useSessionClient } from "../useSessionClient";
 import { useShowQueryError } from "../errors/useShowQueryError";
 import { useShowMutationError } from "../errors/useShowMutationError";
 import { StepContentProps } from "./common";
+import { usePrerequisites } from "./usePrerequisites";
 
 export type Props = StepContentProps & {
   userClient: ConnectedClient;
   registerSpender: boolean;
   registerDelegation: boolean;
+  sessionAddress?: Hex;
 };
 
 export function Session({ isActive, isExpanded, userClient, registerSpender, registerDelegation }: Props) {
   const sessionClient = useShowQueryError(useSessionClient(userClient.account.address));
   const setup = useShowMutationError(useSetupSession({ userClient }));
   const hasSession = !registerDelegation && !registerDelegation;
+  const { data: prerequisites } = usePrerequisites(userClient.account.address);
+  const { hasAllowance, hasGasBalance, hasQuarryGasBalance } = prerequisites ?? {};
 
   useEffect(() => {
     // There seems to be a tanstack-query bug(?) where multiple simultaneous renders loses
     // state between the two mutations. They're not treated as shared state but rather
     // individual mutations, even though the keys match. And the one we want the status of
     // seems to stay pending. This is sorta resolved by triggering this after a timeout.
     const timer = setTimeout(() => {
-      if (isActive && setup.status === "idle" && sessionClient.data && !hasSession) {
+      if (
+        isActive &&
+        setup.status === "idle" &&
+        sessionClient.data &&
+        !hasSession &&
+        (hasAllowance || hasGasBalance || hasQuarryGasBalance)
+      ) {
         setup.mutate({
           sessionClient: sessionClient.data,
           registerSpender,
@@ -33,7 +44,17 @@ export function Session({ isActive, isExpanded, userClient, registerSpender, reg
       }
     });
     return () => clearTimeout(timer);
-  }, [hasSession, isActive, registerDelegation, registerSpender, sessionClient, setup]);
+  }, [
+    hasSession,
+    isActive,
+    registerDelegation,
+    registerSpender,
+    sessionClient,
+    setup,
+    hasAllowance,
+    hasGasBalance,
+    hasQuarryGasBalance,
+  ]);
 
   return (
     <div className="flex flex-col gap-4">

packages/entrykit/src/onboarding/usePrerequisites.ts

@@ -54,14 +54,14 @@ export function getPrequisitesQueryOptions({
             const hasAllowance = allowance == null || allowance > 0n;
             const isSpender = spender == null ? true : spender;
             const hasGasBalance = sessionBalance == null || sessionBalance.value > 0n;
-            const hasQuarryBalance = quarryBalance == null || quarryBalance > 0n;
+            const hasQuarryGasBalance = quarryBalance == null || quarryBalance > 0n;
 
             return {
               sessionAddress,
               hasAllowance,
               isSpender,
               hasGasBalance,
-              hasQuarryBalance,
+              hasQuarryGasBalance,
               hasDelegation,
               // we intentionally don't enforce an allowance/gas balance here
               complete: isSpender && hasDelegation,