perf(wallet): bound and parallelize whoami status queries

src/commands/identity.ts

@@ -408,20 +408,22 @@ export async function currentWhoamiOutput(options: {
         )
       : undefined);
   const token = key?.limits[0]?.token ?? tokenAddress(selectedChain);
-  const balance = await tokenBalance({
-    token,
-    walletAddress: options.walletAddress,
-    network: options.network,
-  });
-  const balances = await walletBalances({
+  const [balance, assets, sessions] = await Promise.all([
+    tokenBalance({
+      token,
+      walletAddress: options.walletAddress,
+      network: options.network,
+    }),
+    fetchWalletAssets({
+      chain: selectedChain,
+      walletAddress: options.walletAddress,
+    }),
+    sessionStats({ token, walletAddress: options.walletAddress }),
+  ]);
+  const balances = walletBalances({
     accessKey: key,
-    chain: selectedChain,
+    assets,
     fallback: balance,
-    walletAddress: options.walletAddress,
-  });
-  const sessions = await sessionStats({
-    token: balance?.token ?? token,
-    walletAddress: options.walletAddress,
   });
   return {
     ready: Boolean(options.walletAddress && paymentKey && balance),
@@ -568,19 +570,12 @@ type WalletAsset = {
   verified: boolean;
 };
 
-async function walletBalances(options: {
+function walletBalances(options: {
   accessKey: WalletState["accessKeys"][number] | undefined;
-  chain: number;
+  assets: WalletAsset[];
   fallback: TokenBalance | null;
-  walletAddress: string | null;
 }) {
-  if (!options.walletAddress) return [];
-
-  const assets = await fetchWalletAssets({
-    chain: options.chain,
-    walletAddress: options.walletAddress,
-  });
-  const balances = assets.map((asset) => {
+  const balances = options.assets.map((asset) => {
     const limit = options.accessKey?.limits.find(
       (candidate) => candidate.token.toLowerCase() === asset.address.toLowerCase(),
     );
@@ -616,17 +611,21 @@ async function walletBalances(options: {
   return balances;
 }
 
+// Status queries should return promptly even when an upstream service stalls.
+const statusQueryTimeout = 3_000;
+
 async function fetchWalletAssets(options: {
   chain: number;
-  walletAddress: string;
+  walletAddress: string | null;
 }): Promise<WalletAsset[]> {
+  if (!options.walletAddress) return [];
   const url = new URL("/api/assets", appUrl);
   url.searchParams.set("address", options.walletAddress);
   url.searchParams.set("chainId", String(options.chain));
   url.searchParams.set("fresh", "true");
 
   try {
-    const response = await fetch(url);
+    const response = await fetch(url, { signal: AbortSignal.timeout(statusQueryTimeout) });
     if (!response.ok) return [];
     return getArray(await response.json()).flatMap((value) => {
       const asset = getRecord(value);
@@ -686,7 +685,10 @@ async function tokenBalance(options: {
   if (!options.walletAddress) return null;
   const token = options.token ?? tokenAddress(chainId(options.network));
   try {
-    const client = createTempoPublicClient(options.network);
+    const client = createTempoPublicClient(options.network, {
+      timeout: statusQueryTimeout,
+      retryCount: 0,
+    });
     const raw = await client.readContract({
       address: token as Address,
       abi: erc20Abi,

src/shared/network.ts

@@ -1,4 +1,4 @@
-import { createPublicClient, http, type Address } from "viem";
+import { createPublicClient, http, type Address, type HttpTransportConfig } from "viem";
 import { Chain } from "viem/tempo";
 
 import { usageError } from "./errors.js";
@@ -34,11 +34,14 @@ export function tokenAddress(chain: number) {
   return (chain === 42431 ? moderatoToken : usdcToken) as Address;
 }
 
-export function createTempoPublicClient(network: string | undefined) {
+export function createTempoPublicClient(
+  network: string | undefined,
+  options: Pick<HttpTransportConfig, "timeout" | "retryCount"> = {},
+) {
   const chain = chainId(network) === 42431 ? Chain.tempoModerato : Chain.tempo;
   return createPublicClient({
     chain,
-    transport: http(rpcUrl(network)),
+    transport: http(rpcUrl(network), options),
   });
 }
 

test/identity.test.ts

@@ -434,6 +434,7 @@ describe("identity commands", () => {
         `/api/assets?address=${testWallet}&chainId=4217&fresh=true`,
         "https://wallet.tempo.xyz",
       ),
+      { signal: expect.any(AbortSignal) },
     );
     expect("balances" in result ? result.balances : null).toEqual([
       {
@@ -455,6 +456,29 @@ describe("identity commands", () => {
     ]);
   });
 
+  it("starts asset discovery without waiting for the payment-token RPC", async () => {
+    await useTempHome();
+    const balance = Promise.withResolvers<bigint>();
+    mocks.readContract.mockReturnValueOnce(balance.promise);
+
+    const output = currentWhoamiOutput({
+      walletAddress: testWallet,
+      chain: 4217,
+      accessKeys: walletState().accessKeys,
+    });
+    try {
+      expect(mocks.readContract).toHaveBeenCalledOnce();
+      expect(mocks.fetch).toHaveBeenCalledOnce();
+    } finally {
+      balance.resolve(5_000_000n);
+    }
+    expect(await output).toMatchObject({
+      ready: true,
+      balance: { available: "5" },
+      balances: [{ token: usdc.toLowerCase(), balance: "5" }],
+    });
+  });
+
   it("whoami preserves the payment-token balance when asset discovery fails", async () => {
     await useTempHome();
     await writeWalletState(walletState());

test/whoami-timeout.test.ts

@@ -0,0 +1,79 @@
+import { execFile } from "node:child_process";
+import { createServer } from "node:http";
+import { resolve } from "node:path";
+import { promisify } from "node:util";
+
+import { expect, it } from "vitest";
+
+import { useTempHome, walletState, writeWalletState } from "./helpers.js";
+
+it.each(["rpc", "assets", "asset body"] as const)(
+  "whoami returns when the %s response stalls",
+  async (stalled) => {
+    const home = await useTempHome();
+    await writeWalletState(walletState());
+    let rpcCalls = 0;
+    let assetCalls = 0;
+    const server = createServer(async (request, response) => {
+      if (request.method === "GET") {
+        assetCalls++;
+        if (stalled === "assets") return;
+        response.writeHead(200, { "Content-Type": "application/json" });
+        if (stalled === "asset body") {
+          response.write("[");
+          return;
+        }
+        response.end("[]");
+        return;
+      }
+      rpcCalls++;
+      const chunks: Buffer[] = [];
+      for await (const chunk of request) chunks.push(Buffer.from(chunk));
+      const rpc = JSON.parse(Buffer.concat(chunks).toString()) as { id: number };
+      if (stalled === "rpc") return;
+      response.writeHead(200, { "Content-Type": "application/json" });
+      response.end(
+        JSON.stringify({
+          jsonrpc: "2.0",
+          id: rpc.id,
+          result: `0x${5_000_000n.toString(16).padStart(64, "0")}`,
+        }),
+      );
+    });
+    await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
+    try {
+      const address = server.address();
+      if (!address || typeof address === "string") throw new Error("Missing server address");
+      const url = `http://127.0.0.1:${address.port}`;
+      const { stdout } = await promisify(execFile)(
+        process.execPath,
+        ["--import", "tsx", "src/cli.ts", "whoami", "--json"],
+        {
+          cwd: resolve(import.meta.dirname, ".."),
+          env: { ...process.env, HOME: home, TEMPO_RPC_URL: url, TEMPO_AUTH_URL: url },
+          // Allow startup overhead while rejecting the old 40-second RPC retry path.
+          timeout: 10_000,
+        },
+      );
+      const output = JSON.parse(stdout);
+      expect(rpcCalls).toBe(1);
+      expect(assetCalls).toBe(1);
+      if (stalled === "rpc") {
+        expect(output).toMatchObject({
+          ready: false,
+          balance: { available: null, error: { code: "E_RPC" } },
+        });
+      } else {
+        expect(output).toMatchObject({
+          ready: true,
+          balance: { available: "5" },
+          balances: [{ balance: "5" }],
+        });
+      }
+    } finally {
+      server.closeAllConnections();
+      await new Promise<void>((resolve) => server.close(() => resolve()));
+    }
+  },
+  15_000,
+);