feat: better balances fetching

Author: vacekj

package.json

@@ -23,32 +23,42 @@
     "test": "vitest run",
     "generate-pindexer-schema": "tsx ./scripts/generate-pindexer-schema.ts"
   },
+  "pnpm": {
+    "overrides": {
+      "cosmos-kit": "2.23.9"
+    }
+  },
   "dependencies": {
     "@bufbuild/protobuf": "^1.10.0",
-    "@chain-registry/types": "^0.50.80",
+    "@chain-registry/types": "^0.50.82",
     "@chakra-ui/icons": "^2.2.4",
     "@chakra-ui/react": "^2.10.4",
     "@chakra-ui/styled-system": "^2.12.0",
     "@connectrpc/connect": "^1.6.1",
     "@connectrpc/connect-web": "^1.6.1",
     "@cosmjs/proto-signing": "^0.33.0",
-    "@cosmjs/stargate": "^0.33.0",
+    "@cosmjs/stargate": "0.33.0",
+    "@cosmos-kit/core": "^2.14.0",
+    "@cosmos-kit/keplr": "^2.14.7",
+    "@cosmos-kit/leap": "^2.14.7",
+    "@cosmos-kit/react": "^2.21.7",
     "@emotion/react": "^11.13.5",
     "@emotion/styled": "^11.13.5",
     "@formkit/auto-animate": "^0.8.2",
     "@grpc/proto-loader": "^0.7.13",
+    "@interchain-ui/react": "^1.26.2",
     "@keplr-wallet/types": "^0.12.190",
-    "@penumbra-labs/registry": "^12.1.1",
-    "@penumbra-zone/bech32m": "^13.0.0",
-    "@penumbra-zone/client": "^24.0.0",
-    "@penumbra-zone/crypto-web": "^36.0.0",
-    "@penumbra-zone/getters": "^23.0.0",
-    "@penumbra-zone/perspective": "^47.0.0",
-    "@penumbra-zone/protobuf": "^7.2.0",
+    "@penumbra-labs/registry": "^12.4.0",
+    "@penumbra-zone/bech32m": "^14.0.0",
+    "@penumbra-zone/client": "^25.0.0",
+    "@penumbra-zone/crypto-web": "^38.0.0",
+    "@penumbra-zone/getters": "^24.1.0",
+    "@penumbra-zone/perspective": "^49.0.0",
+    "@penumbra-zone/protobuf": "^8.0.0",
     "@penumbra-zone/transport-dom": "^7.5.0",
-    "@penumbra-zone/types": "29.1.0",
-    "@penumbra-zone/ui": "^13.11.0",
-    "@penumbra-zone/wasm": "40.0.0",
+    "@penumbra-zone/types": "31.0.0",
+    "@penumbra-zone/ui": "^13.13.0",
+    "@penumbra-zone/wasm": "42.0.0",
     "@radix-ui/react-icons": "^1.3.2",
     "@rehooks/component-size": "^1.0.3",
     "@tanstack/react-query": "^5.62.3",
@@ -57,14 +67,15 @@
     "@tsconfig/vite-react": "^3.4.0",
     "bech32": "^2.0.0",
     "bignumber.js": "^9.1.2",
-    "chain-registry": "^1.69.134",
+    "chain-registry": "^1.69.136",
     "chart.js": "^4.4.7",
     "chartjs-adapter-date-fns": "^3.0.0",
     "chartjs-chart-financial": "^0.2.1",
     "chartjs-plugin-annotation": "^3.1.0",
     "chartjs-plugin-zoom": "^2.2.0",
     "clsx": "^2.1.1",
     "configs": "github:prax-wallet/configs#main",
+    "cosmos-kit": "^2.23.9",
     "date-fns": "^4.1.0",
     "echarts": "^5.5.1",
     "echarts-for-react": "^3.0.2",
@@ -78,6 +89,7 @@
     "mobx": "^6.13.5",
     "mobx-react-lite": "^4.0.7",
     "next": "^14.2.20",
+    "osmo-query": "^16.14.0",
     "pg": "^8.13.1",
     "prettier": "^3.4.2",
     "react": "^18.3.1",

pnpm-lock.yaml

@@ -0,0 +1,79 @@
+import { ChainProvider } from '@cosmos-kit/react';
+import { assets, chains } from 'chain-registry';
+import { SignerOptions, wallets } from 'cosmos-kit';
+import { ReactNode, useMemo } from 'react';
+import { Chain, Registry as PenumbraRegistry } from '@penumbra-labs/registry';
+
+import { GeneratedType, Registry } from '@cosmjs/proto-signing';
+import { AminoTypes } from '@cosmjs/stargate';
+import {
+  cosmosAminoConverters,
+  cosmosProtoRegistry,
+  cosmwasmAminoConverters,
+  cosmwasmProtoRegistry,
+  ibcAminoConverters,
+  ibcProtoRegistry,
+  osmosisAminoConverters,
+  osmosisProtoRegistry,
+} from 'osmo-query';
+
+// @ts-expect-error type error, but it works.
+import '@interchain-ui/react/styles';
+
+const protoRegistry: readonly [string, GeneratedType][] = [
+  ...cosmosProtoRegistry,
+  ...cosmwasmProtoRegistry,
+  ...ibcProtoRegistry,
+  ...osmosisProtoRegistry,
+];
+
+const aminoConverters = {
+  ...cosmosAminoConverters,
+  ...cosmwasmAminoConverters,
+  ...ibcAminoConverters,
+  ...osmosisAminoConverters,
+};
+
+const registry = new Registry(protoRegistry);
+const aminoTypes = new AminoTypes(aminoConverters);
+
+const signerOptions: SignerOptions = {
+  // @ts-expect-error type error, but it works.
+  signingStargate: () => {
+    return {
+      aminoTypes,
+      registry,
+    };
+  },
+};
+
+interface IbcProviderProps {
+  registry: PenumbraRegistry;
+  children: ReactNode;
+}
+
+export const IbcChainProvider = ({ registry, children }: IbcProviderProps) => {
+  const chainsToDisplay = useMemo(
+    () => chainsInPenumbraRegistry(registry.ibcConnections),
+    [registry],
+  );
+
+  return (
+    <ChainProvider
+      chains={chainsToDisplay}
+      assetLists={assets}
+      // Not using mobile wallets as WalletConnect is a centralized service that requires an account
+      wallets={wallets.extension}
+      signerOptions={signerOptions}
+      modalTheme={{ defaultTheme: 'light' }}
+      logLevel={'NONE'}
+    >
+      {children}
+    </ChainProvider>
+  );
+};
+
+// Searches cosmos registry for chains that have ibc connections to Penumbra
+export const chainsInPenumbraRegistry = (ibcConnections: Chain[]) => {
+  return chains.filter(c => ibcConnections.some(i => c.chain_id === i.chainId));
+};

src/features/cosmos/chain-provider.tsx

@@ -0,0 +1,56 @@
+import { observer } from 'mobx-react-lite';
+import { Button, ButtonProps } from '@penumbra-zone/ui/Button';
+import dynamic from 'next/dynamic';
+import { useChains } from '@cosmos-kit/react';
+import { Wallet2 } from 'lucide-react';
+import { Density } from '@penumbra-zone/ui/Density';
+import { chainsInPenumbraRegistry } from '@/features/cosmos/chain-provider.tsx';
+import { useRegistry } from '@/shared/api/registry.ts';
+import { useBalances } from '@/features/cosmos/use-augmented-balances.ts';
+
+interface CosmosConnectButtonProps {
+  actionType?: ButtonProps['actionType'];
+  variant?: 'default' | 'minimal';
+}
+
+const CosmosConnectButtonInner = observer(
+  ({ actionType = 'accent', variant = 'default' }: CosmosConnectButtonProps) => {
+    const { data: registry } = useRegistry();
+    const penumbraIbcChains = chainsInPenumbraRegistry(registry?.ibcConnections ?? []).map(
+      c => c.chain_name,
+    );
+    const chains = useChains(penumbraIbcChains);
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion -- Osmosis is always available
+    const { address, disconnect, openView, isWalletConnected } = chains['osmosis']!;
+
+    const handleConnect = () => {
+      openView();
+    };
+
+    const { balances } = useBalances();
+
+    return (
+      <Density variant={variant === 'default' ? 'sparse' : 'compact'}>
+        {isWalletConnected && address ? (
+          <>
+            <Button actionType={actionType} onClick={() => void disconnect()}>
+              {`${address.slice(0, 8)}...${address.slice(-4)}`}
+            </Button>
+            <pre className={'text-white overflow-scroll'}>
+              {balances.map(b => JSON.stringify(b, undefined, 4))}
+            </pre>
+          </>
+        ) : (
+          <Button icon={Wallet2} actionType={actionType} onClick={handleConnect}>
+            Connect Cosmos Wallet
+          </Button>
+        )}
+      </Density>
+    );
+  },
+);
+
+// Export a dynamic component to prevent SSR issues with window object
+export const CosmosConnectButton = dynamic(() => Promise.resolve(CosmosConnectButtonInner), {
+  ssr: false,
+});

src/features/cosmos/cosmos-connect-button.tsx

@@ -1,5 +1,3 @@
-import type { Asset, Chain } from '@chain-registry/types';
-
 export interface Balance {
   denom: string;
   amount: string;
@@ -8,13 +6,6 @@ export interface Balance {
   displayAmount?: string;
 }
 
-export interface ChainConfig {
-  chainId: string;
-  restEndpoint: string;
-  chain: Chain & { bech32_prefix: string };
-  assets: Asset[];
-}
-
 export interface BankBalancesResponse {
   balances: {
     denom: string;

src/features/cosmos/get-cosmos-balances.tsx

@@ -0,0 +1,65 @@
+import { useState, useEffect } from 'react';
+import { fetchChainBalances } from './utils/fetch-balances';
+import { useWallet } from '@cosmos-kit/react';
+import { Balance } from '@/features/cosmos/types.ts';
+
+/**
+ * Hook that augments Cosmos balances with Penumbra-specific data
+ *
+ * This scans balances across all supported chains and adds:
+ * - Display information (symbol, denom, amount)
+ * - Asset icons
+ * - Asset type information
+ * - A flag indicating if an asset is a Penumbra asset (can be transferred via IBC)
+ */
+export const useBalances = () => {
+  const [balances, setBalances] = useState<Balance[]>([]);
+  const [isLoading, setIsLoading] = useState(false);
+  const [error, setError] = useState<string | null>(null);
+  const { chainWallets, status } = useWallet();
+
+  useEffect(() => {
+    const scanAllChains = async () => {
+      try {
+        setIsLoading(true);
+        setError(null);
+
+        const allBalances: Balance[] = [];
+
+        // For each chain, try to fetch balances
+        await Promise.all(
+          chainWallets.map(async chain => {
+            try {
+              // Get the address for this chain
+              const address = chain.address;
+
+              if (!address) {
+                return;
+              }
+
+              // Fetch balances for this chain
+              const chainBalances = await fetchChainBalances(address, chain);
+              allBalances.push(...chainBalances);
+            } catch (err) {
+              console.warn(`Error fetching balances for chain ${chain.chainId}:`, err);
+            }
+          }),
+        );
+
+        setBalances(allBalances);
+      } catch (err) {
+        console.error('Error scanning balances:', err);
+        setError(err instanceof Error ? err.message : 'Unknown error');
+      } finally {
+        setIsLoading(false);
+      }
+    };
+
+    void scanAllChains();
+  }, [status]);
+  return {
+    balances,
+    isLoading: isLoading,
+    error: error,
+  };
+};