Add option to stop and remove all liquidity in one tx and improve withdraw UX

apps/cow-amm-deployer/src/app/[userId]/amms/[id]/(components)/DisableAmmButton.tsx

@@ -2,32 +2,86 @@
 
 import { toast } from "@bleu/ui";
 import { StopIcon } from "@radix-ui/react-icons";
-import React from "react";
+import React, { useEffect } from "react";
+import { parseUnits } from "viem";
 import { useAccount } from "wagmi";
 
 import { Button } from "#/components/Button";
+import { Checkbox } from "#/components/Checkbox";
+import { Dialog } from "#/components/Dialog";
 import { useManagedTransaction } from "#/hooks/tx-manager/useManagedTransaction";
 import { ICowAmm } from "#/lib/fetchAmmData";
-import { DisableCoWAMMArgs, TRANSACTION_TYPES } from "#/lib/transactionFactory";
+import {
+  AllTransactionArgs,
+  TRANSACTION_TYPES,
+  WithdrawCoWAMMArgs,
+} from "#/lib/transactionFactory";
 import { ChainId } from "#/utils/chainsPublicClients";
 
 export function DisableAmmButton({ ammData }: { ammData: ICowAmm }) {
+  const [isOpen, setIsOpen] = React.useState(false);
+
+  return (
+    <Dialog
+      title="Disable"
+      content={
+        <DisableTradingDialogContent ammData={ammData} setIsOpen={setIsOpen} />
+      }
+      isOpen={isOpen}
+      onClose={() => setIsOpen(false)}
+      setIsOpen={setIsOpen}
+    >
+      <Button
+        className="flex items-center gap-1 py-3 px-6"
+        variant="destructive"
+        onClick={() => {
+          setIsOpen(true);
+        }}
+      >
+        <StopIcon />
+        Disable trading
+      </Button>
+    </Dialog>
+  );
+}
+
+function DisableTradingDialogContent({
+  ammData,
+  setIsOpen,
+}: {
+  ammData: ICowAmm;
+  setIsOpen: (isOpen: boolean) => void;
+}) {
+  const [withdrawFunds, setWithdrawFunds] = React.useState(false);
   const { chainId } = useAccount();
 
   const { writeContract, writeContractWithSafe, status, isWalletContract } =
     useManagedTransaction();
 
   function onDisableAMM() {
-    const txArgs = {
-      type: TRANSACTION_TYPES.DISABLE_COW_AMM,
-      amm: ammData.order.owner,
-      chainId: chainId as ChainId,
-      hash: ammData.order.hash,
-      version: ammData.version,
-    } as DisableCoWAMMArgs;
+    const txArgs = [
+      {
+        type: TRANSACTION_TYPES.DISABLE_COW_AMM,
+        amm: ammData.order.owner,
+        chainId: chainId as ChainId,
+        hash: ammData.order.hash,
+        version: ammData.version,
+      },
+    ] as AllTransactionArgs[];
+
+    if (withdrawFunds) {
+      txArgs.push({
+        type: TRANSACTION_TYPES.WITHDRAW_COW_AMM,
+        amm: ammData.order.owner,
+        amount0: parseUnits(ammData.token0.balance, ammData.token0.decimals),
+        amount1: parseUnits(ammData.token1.balance, ammData.token1.decimals),
+        chainId: chainId as ChainId,
+      } as WithdrawCoWAMMArgs);
+    }
+
     try {
       if (isWalletContract) {
-        writeContractWithSafe([txArgs]);
+        writeContractWithSafe(txArgs);
       } else {
         // TODO: this will need to be refactored once we have EOAs
         // @ts-ignore
@@ -42,16 +96,38 @@ export function DisableAmmButton({ ammData }: { ammData: ICowAmm }) {
     }
   }
 
+  useEffect(() => {
+    if (status === "confirmed") {
+      setIsOpen(false);
+    }
+  }, [status]);
+
   return (
-    <Button
-      className="flex items-center gap-1 py-3 px-6"
-      variant="destructive"
-      onClick={onDisableAMM}
-      loading={!["final", "idle", "confirmed", "error"].includes(status || "")}
-      loadingText="Confirming..."
-    >
-      <StopIcon />
-      Disable trading
-    </Button>
+    <div className="flex flex-col gap-2 text-background bg-foreground">
+      <span>
+        This action will halt the automatic rebalancing of your AMM over time.
+        You retain the ability to withdraw or deposit funds as needed. If
+        desired, you have the option to withdraw all your funds in one single
+        transaction.
+      </span>
+      <Checkbox
+        label="Withdraw all funds"
+        onChange={() => setWithdrawFunds(!withdrawFunds)}
+        checked={withdrawFunds}
+        id="withdraw-funds-checkbox"
+      />
+      <Button
+        className="flex items-center gap-1 py-3 px-6"
+        variant="destructive"
+        onClick={onDisableAMM}
+        loading={
+          !["final", "idle", "confirmed", "error"].includes(status || "")
+        }
+        loadingText="Confirming..."
+      >
+        <StopIcon />
+        Disable trading
+      </Button>
+    </div>
   );
 }

apps/cow-amm-deployer/src/app/[userId]/amms/[id]/(components)/EditAMMForm.tsx

@@ -76,7 +76,7 @@ export function EditAMMForm({ ammData }: { ammData: ICowAmm }) {
 
   return (
     // @ts-ignore
-    <Form {...form} onSubmit={onSubmit} className="flex flex-col gap-y-3">
+    <Form {...form} onSubmit={onSubmit} className="flex flex-col gap-y-5">
       <div className="flex flex-col w-full">
         <span className="mb-2 h-5 block text-sm">Token Pair</span>
         <div className="flex h-fit gap-x-7">

apps/cow-amm-deployer/src/components/DepositForm.tsx

@@ -73,7 +73,7 @@ export function DepositForm({
 
   return (
     // @ts-ignore
-    <Form {...form} onSubmit={onSubmit} className="flex flex-col gap-y-3">
+    <Form {...form} onSubmit={onSubmit} className="flex flex-col gap-y-5">
       <div className="flex gap-x-2 w-full items-start justify-between">
         <div className="w-1/3">
           <TokenInfo token={ammData.token0} showBalance={false} />

apps/cow-amm-deployer/src/components/Dialog.tsx

@@ -60,7 +60,7 @@ export function Dialog({
           </div>
           <DialogPrimitive.Close asChild>
             <button
-              className="absolute right-[10px] top-[10px] inline-flex size-[30px] items-center justify-center text-sand12 hover:font-black focus:outline-none"
+              className="absolute right-[10px] top-[10px] inline-flex size-[30px] items-center justify-center text-background hover:font-black focus:outline-none"
               aria-label="Close"
             >
               <Cross2Icon />

apps/cow-amm-deployer/src/components/WithdrawForm.tsx

@@ -24,7 +24,7 @@ export function WithdrawForm({ ammData }: { ammData: ICowAmm }) {
     // @ts-ignore
     resolver: zodResolver(ammWithdrawSchema),
     defaultValues: {
-      withdrawPct: 50,
+      withdrawPct: 0,
     },
   });
   const { chainId } = useAccount();
@@ -33,14 +33,22 @@ export function WithdrawForm({ ammData }: { ammData: ICowAmm }) {
     useManagedTransaction();
 
   const onSubmit = async (data: typeof ammWithdrawSchema._type) => {
-    const amount0 = parseUnits(
-      String((Number(ammData.token0.balance) * data.withdrawPct) / 100),
-      ammData.token0.decimals,
-    );
-    const amount1 = parseUnits(
-      String((Number(ammData.token1.balance) * data.withdrawPct) / 100),
-      ammData.token1.decimals,
-    );
+    let amount0 = BigInt(0);
+    let amount1 = BigInt(0);
+    if (data.withdrawPct === 100) {
+      // avoid floating point arithmetic
+      amount0 = parseUnits(ammData.token0.balance, ammData.token0.decimals);
+      amount1 = parseUnits(ammData.token1.balance, ammData.token1.decimals);
+    } else {
+      amount0 = parseUnits(
+        String((Number(ammData.token0.balance) * data.withdrawPct) / 100),
+        ammData.token0.decimals,
+      );
+      amount1 = parseUnits(
+        String((Number(ammData.token1.balance) * data.withdrawPct) / 100),
+        ammData.token1.decimals,
+      );
+    }
     const txArgs = {
       type: TRANSACTION_TYPES.WITHDRAW_COW_AMM,
       amm: ammData.order.owner,
@@ -66,17 +74,33 @@ export function WithdrawForm({ ammData }: { ammData: ICowAmm }) {
 
   const {
     control,
+    setValue,
     formState: { isSubmitting },
   } = form;
 
   const withdrawPct = useWatch({ control, name: "withdrawPct" });
 
   return (
-    <Form {...form} onSubmit={onSubmit} className="flex flex-col gap-y-3">
+    <Form {...form} onSubmit={onSubmit} className="flex flex-col gap-y-5">
       <div className="flex flex-col w-full">
-        <span className="mb-2 h-5 block">
-          Withdraw percentage: {withdrawPct}%
-        </span>
+        <div className="flex justify-between mb-2 items-center">
+          <span className="block text-xl bg-primary p-2">{withdrawPct}%</span>
+          <div className="flex justify-between w-1/2">
+            {[25, 50, 75, 100].map((pct) => (
+              <Button
+                key={pct}
+                type="button"
+                className="px-2"
+                variant="ghost"
+                disabled={isSubmitting || !pct}
+                onClick={() => setValue("withdrawPct", pct)}
+              >
+                {pct === 100 ? "Max" : `${pct}%`}
+              </Button>
+            ))}
+          </div>
+        </div>
+
         <Controller
           name="withdrawPct"
           control={control}
@@ -85,8 +109,8 @@ export function WithdrawForm({ ammData }: { ammData: ICowAmm }) {
               className="relative flex items-center select-none touch-none w-full h-5"
               max={100}
               step={0.1}
-              min={0.1}
-              onValueChange={field.onChange}
+              min={0}
+              onValueChange={(value) => field.onChange(value[0])}
               value={[field.value]}
             >
               <Slider.Track className="relative grow rounded-full h-[3px]">
@@ -112,10 +136,16 @@ export function WithdrawForm({ ammData }: { ammData: ICowAmm }) {
       </div>
       <Button
         variant="highlight"
+        disabled={
+          !withdrawPct ||
+          (ammData.token0.balance === "0" && ammData.token1.balance === "0")
+        }
         type="submit"
         className="w-full"
-        disabled={isSubmitting || (status !== "final" && status !== "idle")}
-        loading={isSubmitting || (status !== "final" && status !== "idle")}
+        loading={
+          isSubmitting ||
+          !["final", "idle", "confirmed", "error"].includes(status || "")
+        }
       >
         Withdraw ${formatNumber((ammData.totalUsdValue * withdrawPct) / 100, 4)}
       </Button>