From 94eecd86e22c3af08c940295b18e71dc3ecdd212 Mon Sep 17 00:00:00 2001
From: Dowland Aiello <dowlandaiello@gmail.com>
Date: Thu, 18 Jul 2024 15:24:11 -0700
Subject: [PATCH] Fix recovery.

---
 src/strategies/util.py | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/src/strategies/util.py b/src/strategies/util.py
index b5c695207..da7b523d1 100644
--- a/src/strategies/util.py
+++ b/src/strategies/util.py
@@ -321,7 +321,6 @@ async def recover_funds(
         [ctx.cli_args["base_denom"], curr_leg.in_asset(), curr_leg.backend.chain_id],
     )
 
-    # TODO: Check this
     route = route[: route.index(curr_leg) - 1]
     backtracked = list(
         reversed([Leg(leg.out_asset, leg.in_asset, leg.backend) for leg in route])
@@ -338,8 +337,21 @@ async def recover_funds(
     if not balance_resp:
         raise ValueError(f"Couldn't get balance for asset {curr_leg.in_asset()}.")
 
-    # TODO: Trade as much of tokens as possible
-    resp = await quantities_for_route_profit(balance_resp, backtracked, r, ctx)
+    if curr_leg.backend.chain_id != backtracked[0].backend.chain_id:
+        to_transfer = min(balance_resp, r.quantities[-2])
+
+        await transfer(
+            r,
+            curr_leg.in_asset(),
+            route[-1].out_asset(),
+            curr_leg,
+            ctx,
+            to_swap,
+        )
+
+    resp = await quantities_for_route_profit(
+        balance_resp, backtracked, r, ctx, seek_profit=False
+    )
 
     if not resp:
         raise ValueError("Couldn't get execution plan.")
@@ -508,6 +520,7 @@ async def quantities_for_route_profit(
     route: list[Leg],
     r: Route,
     ctx: Ctx[Any],
+    seek_profit: Optional[bool] = True,
 ) -> tuple[int, list[int]]:
     """
     Calculates what quantities should be used to obtain
@@ -519,7 +532,9 @@ async def quantities_for_route_profit(
 
     quantities: list[int] = [starting_amount]
 
-    while quantities[-1] - quantities[0] <= 0 or len(quantities) <= len(route):
+    while (seek_profit and quantities[-1] - quantities[0] <= 0) or len(
+        quantities
+    ) <= len(route):
         ctx.log_route(r, "info", "Route has possible execution plan: %s", [quantities])
 
         if starting_amount < DENOM_QUANTITY_ABORT_ARB: