Portfolio Onboarding Section

package.json

@@ -23,17 +23,32 @@
     "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.93",
     "@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",
+    "@cosmos-kit/core": "^2.15.5",
+    "@cosmos-kit/cosmostation": "^2.14.6",
+    "@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.199",
     "@penumbra-labs/registry": "^12.4.0",
     "@penumbra-zone/bech32m": "^15.0.0",
     "@penumbra-zone/client": "^26.0.0",
@@ -53,13 +68,15 @@
     "@tsconfig/vite-react": "^3.4.0",
     "bech32": "^2.0.0",
     "bignumber.js": "^9.1.2",
+    "chain-registry": "^1.69.151",
     "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",
@@ -73,6 +90,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",

src/features/cosmos/chain-provider.tsx

@@ -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/cosmos-connect-button.tsx

@@ -0,0 +1,50 @@
+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';
+
+interface CosmosConnectButtonProps {
+  actionType?: ButtonProps['actionType'];
+  variant?: 'default' | 'minimal';
+  children?: React.ReactNode;
+}
+
+const CosmosConnectButtonInner = observer(
+  ({ actionType = 'accent', variant = 'default', children }: 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'] ??
+      chains['osmosistestnet'])!;
+
+    const handleConnect = () => {
+      openView();
+    };
+
+    return (
+      <Density variant={variant === 'default' ? 'sparse' : 'compact'}>
+        {isWalletConnected && address ? (
+          <Button actionType={actionType} onClick={() => void disconnect()}>
+            {`${address.slice(0, 8)}...${address.slice(-4)}`}
+          </Button>
+        ) : (
+          <Button icon={Wallet2} actionType={actionType} onClick={handleConnect}>
+            {children ?? '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/use-augmented-balances.tsx

@@ -0,0 +1,111 @@
+import { useWallet } from '@cosmos-kit/react';
+import { useQuery } from '@tanstack/react-query';
+import { ChainWalletBase, WalletStatus } from '@cosmos-kit/core';
+import { useRegistry } from '@/shared/api/registry';
+import { Chain } from '@penumbra-labs/registry';
+import { sha256HashStr } from '@penumbra-zone/crypto-web/sha256';
+import { Asset } from '@chain-registry/types';
+import { assets as cosmosAssetList } from 'chain-registry';
+import { Coin, StargateClient } from '@cosmjs/stargate';
+
+export const fetchChainBalances = async (
+  address: string,
+  chain: ChainWalletBase,
+): Promise<readonly Coin[]> => {
+  const endpoint = await chain.getRpcEndpoint();
+  const client = await StargateClient.connect(endpoint);
+  return client.getAllBalances(address);
+};
+
+// Searches for corresponding denom in asset registry and returns the metadata
+export const augmentToAsset = (denom: string, chainName: string): Asset => {
+  const match = cosmosAssetList
+    .find(({ chain_name }) => chain_name === chainName)
+    ?.assets.find(asset => asset.base === denom);
+
+  return match ? match : fallbackAsset(denom);
+};
+
+const fallbackAsset = (denom: string): Asset => {
+  return {
+    base: denom,
+    denom_units: [{ denom, exponent: 0 }],
+    display: denom,
+    name: denom,
+    symbol: denom,
+    type_asset: 'sdk.coin',
+  };
+};
+
+const generatePenumbraIbcDenoms = async (chains: Chain[]): Promise<string[]> => {
+  const ibcAddrs: string[] = [];
+  for (const c of chains) {
+    const ibcStr = `transfer/${c.counterpartyChannelId}/upenumbra`;
+    const encoder = new TextEncoder();
+    const encodedString = encoder.encode(ibcStr);
+
+    const hash = await sha256HashStr(encodedString);
+    ibcAddrs.push(`ibc/${hash.toUpperCase()}`);
+  }
+  return ibcAddrs;
+};
+
+export const usePenumbraIbcDenoms = () => {
+  const { data: registry, isLoading: registryIsLoading, error: registryErr } = useRegistry();
+
+  const {
+    data: ibcAddrs,
+    isLoading: ibcAddrsLoading,
+    error: ibcAddrssErr,
+  } = useQuery({
+    queryKey: ['penumbraIbcDenoms', registry],
+    queryFn: async () => generatePenumbraIbcDenoms(registry?.ibcConnections ?? []),
+    enabled: Boolean(registry),
+  });
+
+  return {
+    data: ibcAddrs,
+    isLoading: registryIsLoading || ibcAddrsLoading,
+    error: registryErr ?? ibcAddrssErr,
+  };
+};
+
+export const useBalances = () => {
+  const { chainWallets, status } = useWallet();
+  const { data: registry } = useRegistry();
+
+  const fetchAllChainBalances = async (): Promise<{ asset: Asset; amount: string }[]> => {
+    return [
+      await Promise.all(
+        chainWallets.map(async chain => {
+          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion -- okay
+          const balances = await fetchChainBalances(chain.address!, chain);
+          return balances.map(coin => {
+            return { asset: augmentToAsset(coin.denom, chain.chainName), amount: coin.amount };
+          });
+        }),
+      ),
+    ].flat(2);
+  };
+
+  const result = useQuery({
+    queryKey: [
+      'cosmos-balances',
+      status,
+      chainWallets.map(chain => chain.chainId).join(','),
+      // Use a stable string representation of the registry to avoid unnecessary invalidation
+      registry ? 'registry-available' : 'no-registry',
+    ],
+    queryFn: fetchAllChainBalances,
+    enabled: status === WalletStatus.Connected && chainWallets.length > 0,
+    retry: 3, // Retry failed requests 3 times
+    retryDelay: attemptIndex => Math.min(1000 * 2 ** attemptIndex, 30000), // Exponential backoff starting at 1s
+  });
+
+  return {
+    balances: result.data ?? [],
+    isLoading: result.isLoading && status === WalletStatus.Connected,
+    error: result.error ? String(result.error) : null,
+    refetch: result.refetch,
+  };
+};