Osmosis Integration Tests + Improvements

local-interchaintest/chains/neutron_osmosis_gaia.json

@@ -167,7 +167,7 @@
       "binary": "osmosisd",
       "bech32_prefix": "osmo",
       "docker_image": {
-        "version": "v25.0.4",
+        "version": "v26.0.2",
         "repository": "ghcr.io/strangelove-ventures/heighliner/osmosis"
       },
       "gas_prices": "0.0025%DENOM%",

local-interchaintest/src/main.rs

@@ -183,7 +183,6 @@ fn main() -> Result<(), Box<dyn StdError + Send + Sync>> {
                 .with_test(Box::new(tests::test_unprofitable_arb) as TestFn)
                 .build()?,
         )?
-        /*
         // Test case (astro -> osmo arb):
         //
         // - Astro: untrn-bruhtoken @ 1.5 bruhtoken/untrn
@@ -213,7 +212,7 @@ fn main() -> Result<(), Box<dyn StdError + Send + Sync>> {
                     bruhtoken_osmo.clone(),
                     Pool::Osmosis(
                         OsmosisPoolBuilder::default()
-                            .with_funds(bruhtoken_osmo.clone(), 10000000u128)
+                            .with_funds(bruhtoken_osmo.clone(), 100000000u128)
                             .with_funds(uosmo.clone(), 10000000u128)
                             .with_weight(bruhtoken_osmo.clone(), 1u128)
                             .with_weight(uosmo.clone(), 1u128)
@@ -223,16 +222,16 @@ fn main() -> Result<(), Box<dyn StdError + Send + Sync>> {
                 .with_pool(
                     untrn.clone(),
                     bruhtoken.clone(),
-                    Pool::Astroport(
-                        AstroportPoolBuilder::default()
-                            .with_balance_asset_a(10000000u128)
-                            .with_balance_asset_b(10000000u128)
+                    Pool::Auction(
+                        AuctionPoolBuilder::default()
+                            .with_balance_offer_asset(10000000u128)
+                            .with_price(Decimal::percent(10))
                             .build()?,
                     ),
                 )
                 .with_arbbot()
                 .with_test(Box::new(tests::test_osmo_arb) as TestFn)
                 .build()?,
-        )?*/
+        )?
         .join()
 }

local-interchaintest/src/setup.rs

@@ -641,7 +641,7 @@ pub fn with_arb_bot_output(test: OwnedTestFn) -> TestResult {
     let proc_handle = Arc::new(proc);
     let proc_handle_watcher = proc_handle.clone();
     let (tx_res, rx_res) = mpsc::channel();
-    let mut finished = AtomicBool::new(false);
+    let finished = AtomicBool::new(false);
 
     let test_handle = test.clone();
 

local-interchaintest/tests/transfer_neutron.py

@@ -46,6 +46,8 @@ async def main() -> None:
             [],
             {},
             denoms,
+            {},
+            {},
         )
 
         await transfer_raw(

local-interchaintest/tests/transfer_osmosis.py

@@ -46,6 +46,8 @@ async def main() -> None:
             [],
             {},
             denoms,
+            {},
+            {},
         )
 
         await transfer_raw(

main.py

@@ -13,7 +13,7 @@
 import sys
 from os import path
 import os
-from typing import Any, cast, Optional
+from typing import Any, cast
 from cosmpy.aerial.client import LedgerClient
 from cosmpy.aerial.wallet import LocalWallet
 from src.scheduler import Scheduler, Ctx
@@ -105,12 +105,16 @@ async def main() -> None:
                 f,
             )
 
-    denom_map: Optional[dict[str, list[dict[str, str]]]] = None
+    denom_file: dict[str, Any] = {
+        "denom_map": {},
+        "denom_routes": {},
+        "chain_info": {},
+    }
 
     # If the user has specified a denom map, use that instead of skip
     if args.denom_file is not None and path.isfile(args.denom_file):
         with open(args.denom_file, "r", encoding="utf-8") as f:
-            denom_map = json.load(f)
+            denom_file = json.load(f)
 
     # If the user specified a poolfile, create the poolfile if it is empty
     if args.pool_file is not None and not path.isfile(args.pool_file):
@@ -188,7 +192,9 @@ async def main() -> None:
                 session,
                 [],
                 cast(dict[str, Any], json.load(f)),
-                denom_map,
+                denom_file["denom_map"],
+                denom_file["denom_routes"],
+                denom_file["chain_info"],
             ).recover_history()
             sched = Scheduler(ctx, strategy)
 

src/contracts/route.py

@@ -53,6 +53,7 @@ class Route:
 
     uid: int
     route: list[LegRepr]
+    legs: list[Leg]
     theoretical_profit: int
     expected_profit: int
     realized_profit: Optional[int]
@@ -104,6 +105,7 @@ def load_route(s: str) -> Route:
     return Route(
         loaded["uid"],
         [load_leg_repr(json_leg) for json_leg in loaded["route"]],
+        [],
         loaded["theoretical_profit"],
         loaded["expected_profit"],
         loaded["realized_profit"],

src/scheduler.py

@@ -8,10 +8,18 @@
 from typing import Callable, List, Self, Optional, Awaitable, Any, TypeVar, Generic
 from dataclasses import dataclass
 from cosmpy.aerial.client import LedgerClient
+from cosmpy.crypto.address import Address
 from cosmpy.aerial.wallet import LocalWallet
 from src.contracts.auction import AuctionDirectory, AuctionProvider
 from src.contracts.route import Route, load_route, LegRepr, Status, Leg
 from src.contracts.pool.provider import PoolProvider
+from src.util import (
+    try_multiple_clients,
+    DenomRouteLeg,
+    DenomRouteQuery,
+    ChainInfo,
+    DenomChainInfo,
+)
 import aiohttp
 import grpc
 
@@ -21,6 +29,10 @@
 MAX_ROUTE_HISTORY_LEN = 200000
 
 
+# Length to truncate denoms in balance logs to
+DENOM_BALANCE_PREFIX_MAX_DENOM_LEN = 12
+
+
 TState = TypeVar("TState")
 
 
@@ -41,7 +53,9 @@ class Ctx(Generic[TState]):
     http_session: aiohttp.ClientSession
     order_history: list[Route]
     deployments: dict[str, Any]
-    denom_map: Optional[dict[str, list[dict[str, str]]]]
+    denom_map: dict[str, list[DenomChainInfo]]
+    denom_routes: dict[DenomRouteQuery, list[DenomRouteLeg]]
+    chain_info: dict[str, ChainInfo]
 
     def with_state(self, state: Any) -> Self:
         """
@@ -104,6 +118,7 @@ def queue_route(
                 LegRepr(leg.in_asset(), leg.out_asset(), leg.backend.kind, False)
                 for leg in route
             ],
+            route,
             theoretical_profit,
             expected_profit,
             None,
@@ -133,14 +148,46 @@ def log_route(
         Writes a log to the standard logger and to the log file of a route.
         """
 
-        route.logs.append(f"{log_level.upper()} {fmt_string % tuple(args)}")
+        def asset_balance_prefix(leg: Leg, asset: str) -> Optional[str]:
+            balance_resp_asset = try_multiple_clients(
+                self.clients[leg.backend.chain_id],
+                lambda client: client.query_bank_balance(
+                    Address(
+                        self.wallet.public_key(),
+                        prefix=leg.backend.chain_prefix,
+                    ),
+                    asset,
+                ),
+            )
+
+            if not balance_resp_asset:
+                return None
+
+            return f"balance[{leg.backend.chain_id}]({asset[:DENOM_BALANCE_PREFIX_MAX_DENOM_LEN]}): {balance_resp_asset}"
+
+        def leg_balance_prefixes(leg: Leg) -> list[str]:
+            assets = [leg.in_asset(), leg.out_asset()]
+
+            return [
+                x for x in (asset_balance_prefix(leg, asset) for asset in assets) if x
+            ]
+
+        prefix = " ".join(
+            {
+                prefix
+                for leg_prefixes in [leg_balance_prefixes(leg) for leg in route.legs]
+                for prefix in leg_prefixes
+            }
+        )
+
+        route.logs.append(f"{log_level.upper()} {prefix} {fmt_string % tuple(args)}")
 
         if route.uid >= len(self.order_history) or route.uid < 0:
             return
 
         self.order_history[route.uid] = route
 
-        fmt_string = f"%s- {fmt_string}"
+        fmt_string = f"{prefix} %s- {fmt_string}"
 
         if log_level == "info":
             logger.info(fmt_string, str(route), *args)
@@ -155,6 +202,163 @@ def log_route(
         if log_level == "debug":
             logger.debug(fmt_string, str(route), *args)
 
+    async def query_denom_route(
+        self, query: DenomRouteQuery
+    ) -> Optional[list[DenomRouteLeg]]:
+        if self.denom_routes and query in self.denom_routes:
+            return self.denom_routes[query]
+
+        head = {"accept": "application/json", "content-type": "application/json"}
+
+        async with self.http_session.post(
+            "https://api.skip.money/v2/fungible/route",
+            headers=head,
+            json={
+                "amount_in": "1",
+                "source_asset_denom": query.src_denom,
+                "source_asset_chain_id": query.src_chain,
+                "dest_asset_denom": query.dest_denom,
+                "dest_asset_chain_id": query.dest_chain,
+                "allow_multi_tx": True,
+                "allow_unsafe": False,
+                "bridges": ["IBC"],
+            },
+        ) as resp:
+            if resp.status != 200:
+                return None
+
+            ops = (await resp.json())["operations"]
+
+            # The transfer includes a swap or some other operation
+            # we can't handle
+            if any(("transfer" not in op for op in ops)):
+                return None
+
+            transfer_info = ops[0]["transfer"]
+
+            from_chain_info = await self.query_chain_info(
+                transfer_info["from_chain_id"]
+            )
+            to_chain_info = await self.query_chain_info(transfer_info["to_chain_id"])
+
+            if not from_chain_info or not to_chain_info:
+                return None
+
+            route = [
+                DenomRouteLeg(
+                    src_chain=query.src_chain,
+                    dest_chain=query.dest_chain,
+                    src_denom=query.src_denom,
+                    dest_denom=query.dest_denom,
+                    from_chain=from_chain_info,
+                    to_chain=to_chain_info,
+                    denom_in=transfer_info["denom_in"],
+                    denom_out=transfer_info["denom_out"],
+                    port=transfer_info["port"],
+                    channel=transfer_info["channel"],
+                )
+                for op in ops
+            ]
+
+            self.denom_routes[query] = route
+
+            return route
+
+    async def query_chain_info(
+        self,
+        chain_id: str,
+    ) -> Optional[ChainInfo]:
+        """
+        Gets basic information about a cosmos chain.
+        """
+
+        if chain_id in self.chain_info:
+            return self.chain_info[chain_id]
+
+        head = {"accept": "application/json", "content-type": "application/json"}
+
+        async with self.http_session.get(
+            f"https://api.skip.money/v2/info/chains?chain_ids={chain_id}",
+            headers=head,
+        ) as resp:
+            if resp.status != 200:
+                return None
+
+            chains = (await resp.json())["chains"]
+
+            if len(chains) == 0:
+                return None
+
+            chain = chains[0]
+
+            chain_info = ChainInfo(
+                chain_name=chain["chain_name"],
+                chain_id=chain["chain_id"],
+                pfm_enabled=chain["pfm_enabled"],
+                supports_memo=chain["supports_memo"],
+                bech32_prefix=chain["bech32_prefix"],
+                fee_asset=chain["fee_assets"][0]["denom"],
+                chain_type=chain["chain_type"],
+                pretty_name=chain["pretty_name"],
+            )
+
+            self.chain_info[chain_id] = chain_info
+
+            return chain_info
+
+    async def query_denom_info_on_chain(
+        self,
+        src_chain: str,
+        src_denom: str,
+        dest_chain: str,
+    ) -> Optional[DenomChainInfo]:
+        """
+        Gets a neutron denom's denom and channel on/to another chain.
+        """
+
+        infos = await self.query_denom_info(src_chain, src_denom)
+
+        return next((info for info in infos if info.dest_chain_id == dest_chain))
+
+    async def query_denom_info(
+        self,
+        src_chain: str,
+        src_denom: str,
+    ) -> list[DenomChainInfo]:
+        """
+        Gets a denom's denom and channel on/to other chains.
+        """
+
+        if src_denom in self.denom_map:
+            return self.denom_map[src_denom]
+
+        head = {"accept": "application/json", "content-type": "application/json"}
+
+        async with self.http_session.post(
+            "https://api.skip.money/v1/fungible/assets_from_source",
+            headers=head,
+            json={
+                "allow_multi_tx": False,
+                "include_cw20_assets": True,
+                "source_asset_denom": src_denom,
+                "source_asset_chain_id": src_chain,
+                "client_id": "timewave-arb-bot",
+            },
+        ) as resp:
+            if resp.status != 200:
+                return []
+
+            dests = (await resp.json())["dest_assets"]
+
+            def chain_info(chain_id: str, info: dict[str, Any]) -> DenomChainInfo:
+                info = info["assets"][0]
+
+                return DenomChainInfo(
+                    src_chain_id=src_chain, denom=info["denom"], dest_chain_id=chain_id
+                )
+
+            return [chain_info(chain_id, info) for chain_id, info in dests.items()]
+
 
 class Scheduler(Generic[TState]):
     """